strims.py

import json
import os
import shutil
import time
import re
from subprocess import Popen


class strims_ffmpeg:
    def __init__(self, dbname="streams.json", hostAddress="0.0.0.0") -> None:
        self.dbname = dbname
        self.hostAddress = hostAddress

    def fromDatabase(self) -> dict:
        with open("streams.json", "r") as json_file:
            return json.load(json_file)

    def son_save(self, infile, data):
        with open(infile, 'w') as outfile:
            json.dump(data, outfile, indent=4)
            outfile.close()
        return

    def check_files(self):
        value = {"default": {"key": "", "link": ""},
                 "streams": [],  "keys": []}

        if not os.path.isfile("{}".format(self.dbname)):
            print(f"Creating empty configuration file: {self.dbname}")
            self.son_save("{}".format(self.dbname), value)

        return

    def startStream(self, streamLink: str, streamKey: str) -> str:
        process = ["ffmpeg"]
        isUrl = re.match(
            "^(http[s]?:\/\/(www\.)?|ftp:\/\/(www\.)?|www\.){1}([0-9A-Za-z-\.@:%_\+~#=]+)+((\.[a-zA-Z]{2,3})+)(\/(.)*)?(\?(.)*)?", streamLink)
        if isUrl:
            process += ["-reconnect", "1", "-reconnect_at_eof",
                        "1", "-reconnect_delay_max", "10"]

        process += ["-re", "-y", "-i", streamLink,
                    "-codec:a", "aac", "-c:v", "copy", "-f", "flv", "-flvflags", "no_duration_filesize",
                    f"rtmp://{self.hostAddress}:1935/live/{streamKey}"]

        return " ".join(process)

    def main(self) -> int:
        if shutil.which("ffmpeg") is None and shutil.which("ffprobe") is None:
            raise Exception("ffmpeg and ffprobe is required")

        strims = self.fromDatabase()
        strimsToOpen = []

        if len(strims["streams"]) == 0:
            try:
                strimsToOpen.append(self.startStream(
                    strims["default"]["link"], strims["default"])["key"])
            except TypeError:
                print(
                    "\033[91m[+] Configuration file is missing data... \033[0m")
        else:
            for s in range(0, len(strims["streams"])):
                strimsToOpen.append(self.startStream(
                    strims["streams"][s], strims["keys"][s]))

        procs = [Popen(i) for i in strimsToOpen]
        for p in procs:
            try:
                p.wait()
            except KeyboardInterrupt:
                try:
                    p.terminate()
                except OSError:
                    pass
                p.wait()
                time.sleep(2)
        return 0


if __name__ == "__main__":
    strims_ffmpeg().check_files()
    exit(strims_ffmpeg(hostAddress="127.0.0.1").main())