rffmpeg/rffmpeg

#!/usr/bin/env python3

# rffmpeg.py - Remote FFMPEG transcoding wrapper
#
#    Copyright (C) 2019-2022  Joshua M. Boniface <joshua@boniface.me>
#                             and Contributors.
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
###############################################################################

import click
import logging
import os
import signal
import sys
import yaml

from contextlib import contextmanager
from grp import getgrnam
from pathlib import Path
from pwd import getpwnam
from re import search
from sqlite3 import connect as sqlite_connect
from subprocess import run
from time import sleep


# Set up the logger
log = logging.getLogger("rffmpeg")


# Open a database connection (context manager)
@contextmanager
def dbconn(config):
    """
    Open a database connection.
    """
    conn = sqlite_connect(config["db_path"])
    conn.execute("PRAGMA foreign_keys = 1")
    cur = conn.cursor()
    yield cur
    conn.commit()
    conn.close()


def fail(msg):
    """
    Output an error message and terminate.
    """
    log.error(msg)
    exit(1)


def load_config():
    """
    Parse the YAML configuration file (either /etc/rffmpeg/rffmpeg.yml or specified by the envvar
    RFFMPEG_CONFIG) and return a standard dictionary of configuration values.
    """

    default_config_file = "/etc/rffmpeg/rffmpeg.yml"
    config_file = os.environ.get("RFFMPEG_CONFIG", default_config_file)

    with open(config_file, "r") as cfgfh:
        try:
            o_config = yaml.load(cfgfh, Loader=yaml.SafeLoader)
        except Exception as e:
            fail("Failed to parse configuration file: {}".format(e))

    config = dict()

    # Parse the base group ("rffmpeg")
    config_base = o_config.get("rffmpeg", dict())
    if not config_base:
        fail("Failed to parse configuration file top level key 'rffmpeg'.")

    # Parse the logging group ("rffmpeg" -> "logging")
    config_logging = config_base.get("logging", dict())
    if config_logging is None:
        config_logging = dict()

    # Parse the directories group ("rffmpeg" -> "directories")
    config_directories = config_base.get("directories", dict())
    if config_directories is None:
        config_directories = dict()

    # Parse the remote group ("rffmpeg" -> "remote")
    config_remote = config_base.get("remote", dict())
    if config_remote is None:
        config_remote = dict()

    # Parse the commands group ("rffmpeg" -> "commands")
    config_commands = config_base.get("commands", dict())
    if config_commands is None:
        config_commands = dict()

    # Parse the keys from the logging group
    config["log_to_file"] = config_logging.get("log_to_file", True)
    config["logfile"] = config_logging.get("logfile", "/var/log/jellyfin/rffmpeg.log")

    # Parse the keys from the state group
    config["state_dir"] = config_directories.get("state", "/var/lib/rffmpeg")
    config["persist_dir"] = config_directories.get("persist", "/run/shm")
    config["dir_owner"] = config_directories.get("owner", "jellyfin")
    config["dir_group"] = config_directories.get("group", "sudo")

    # Parse the keys from the remote group
    config["remote_user"] = config_remote.get("user", "jellyfin")
    config["remote_args"] = config_remote.get(
        "args", ["-i", "/var/lib/jellyfin/.ssh/id_rsa"]
    )
    if config["remote_args"] is None:
        config["remote_args"] = []
    config["persist_time"] = config_remote.get("persist", 300)

    # Parse the keys from the commands group
    config["ssh_command"] = config_commands.get("ssh", "/usr/bin/ssh")
    config["pre_commands"] = config_commands.get("pre", [])
    if config["pre_commands"] is None:
        config["pre_commands"] = []
    config["ffmpeg_command"] = config_commands.get(
        "ffmpeg", "/usr/lib/jellyfin-ffmpeg/ffmpeg"
    )
    config["ffprobe_command"] = config_commands.get(
        "ffprobe", "/usr/lib/jellyfin-ffprobe/ffprobe"
    )
    config["fallback_ffmpeg_command"] = config_commands.get(
        "fallback_ffmpeg", "/usr/lib/jellyfin-ffmpeg/ffmpeg"
    )
    config["fallback_ffprobe_command"] = config_commands.get(
        "fallback_ffprobe", "/usr/lib/jellyfin-ffprobe/ffprobe"
    )

    # Set the database path
    config["db_path"] = config["state_dir"] + "/rffmpeg.db"

    # Set a list of special flags that cause different behaviour
    config["special_flags"] = [
        "-version",
        "-encoders",
        "-decoders",
        "-hwaccels",
        "-filters",
        "-h",
    ]

    # Set the current PID of this process
    config["current_pid"] = os.getpid()

    return config


def cleanup(signum="", frame=""):
    """
    Clean up this processes stored transient data.
    """
    global config, p

    with dbconn(config) as cur:
        cur.execute("DELETE FROM states WHERE process_id = ?", (config["current_pid"],))
        cur.execute(
            "DELETE FROM processes WHERE process_id = ?", (config["current_pid"],)
        )


def generate_ssh_command(config, target_host):
    """
    Generate an SSH command for use.
    """
    ssh_command = list()

    # Add SSH component
    ssh_command.append(config["ssh_command"])
    ssh_command.append("-q")
    ssh_command.append("-t")
    ssh_command.append("-t")

    # Set our connection details
    ssh_command.extend(["-o", "ConnectTimeout=1"])
    ssh_command.extend(["-o", "ConnectionAttempts=1"])
    ssh_command.extend(["-o", "StrictHostKeyChecking=no"])
    ssh_command.extend(["-o", "UserKnownHostsFile=/dev/null"])

    # Use SSH control persistence to keep sessions alive for subsequent commands
    if config["persist_time"] > 0:
        ssh_command.extend(["-o", "ControlMaster=auto"])
        ssh_command.extend(
            ["-o", "ControlPath={}/ssh-%r@%h:%p".format(config["persist_dir"])]
        )
        ssh_command.extend(["-o", "ControlPersist={}".format(config["persist_time"])])

    # Add the remote config args
    for arg in config["remote_args"]:
        if arg:
            ssh_command.append(arg)

    # Add user+host string
    ssh_command.append("{}@{}".format(config["remote_user"], target_host))

    return ssh_command


def run_command(command, stdin, stdout, stderr):
    """
    Execute the command using subprocess.
    """
    log.info("Executing command")
    p = run(
        command,
        shell=False,
        bufsize=0,
        universal_newlines=True,
        stdin=stdin,
        stdout=stdout,
        stderr=stderr,
    )
    return p.returncode


def get_target_host(config):
    """
    Determine an optimal target host via data on currently active processes and states.
    """
    # Select all hosts and active processes from the database
    with dbconn(config) as cur:
        hosts = cur.execute("SELECT * FROM hosts").fetchall()
        processes = cur.execute("SELECT * FROM processes").fetchall()

    # Generate a mapping dictionary of hosts and processes
    host_mappings = dict()
    for host in hosts:
        hid, hostname, weight = host

        # Get the latest state
        with dbconn(config) as cur:
            current_state = cur.execute(
                "SELECT * FROM states WHERE host_id = ? ORDER BY id DESC", (hid,)
            ).fetchone()

        if not current_state:
            current_state = "idle"
        else:
            current_state = current_state[3]

        # Create the mappings entry
        host_mappings[hid] = {
            "hostname": hostname,
            "weight": weight,
            "current_state": current_state,
            "commands": [proc[2] for proc in processes if proc[1] == hid],
        }

    lowest_count = 9999
    target_hid = None
    target_hostname = None
    # For each host in the mapping, let's determine if it is suitable
    for hid, host in host_mappings.items():
        # If it's marked as bad, continue
        if host["current_state"] == "bad":
            continue

        # Try to connect to the host and run a very quick command to determine if it is workable
        if host["hostname"] not in ["localhost", "127.0.0.1"]:
            test_ssh_command = generate_ssh_command(config, host["hostname"])
            test_ffmpeg_command = [config["ffmpeg_command"], "-version"]
            retcode = run_command(
                test_ssh_command + test_ffmpeg_command, None, None, None
            )
            if retcode != 0:
                # Mark the host as bad
                with dbconn(config) as cur:
                    log.warning(
                        "Marking host {} as bad due to retcode {}".format(
                            host["hostname"], retcode
                        )
                    )
                    log.info(
                        "SSH command was: {}".format(
                            " ".join(test_ssh_command + test_ffmpeg_command)
                        )
                    )
                    cur.execute(
                        "INSERT INTO states (host_id, process_id, state) VALUES (?, ?, ?)",
                        (hid, config["current_pid"], "bad"),
                    )
                    continue

        # If the host state is idle, we can use it immediately
        if host["current_state"] == "idle":
            target_hid = hid
            target_hostname = host["hostname"]
            break

        # Get the modified count of the host
        raw_proc_count = len(host["commands"])
        weighted_proc_count = raw_proc_count // host["weight"]

        # If this host is currently the least used, provisionally set it as the target
        if weighted_proc_count < lowest_count:
            lowest_count = weighted_proc_count
            target_hid = hid
            target_hostname = host["hostname"]

    log.info("Found optimal host '{}' (ID '{}')".format(target_hostname, target_hid))
    return target_hid, target_hostname


def run_local_ffmpeg(config, ffmpeg_args):
    """
    Run ffmpeg locally, either because "localhost" is the target host, or because no good target
    host was found by get_target_host().
    """
    rffmpeg_ffmpeg_command = list()

    # Prepare our default stdin/stdout/stderr
    stdin = sys.stdin
    stderr = sys.stderr

    if cmd_name == "ffprobe":
        # If we're in ffprobe mode use that command and sys.stdout as stdout
        rffmpeg_ffmpeg_command.append(config["fallback_ffprobe_command"])
        stdout = sys.stdout
    else:
        # Otherwise, we use stderr as stdout
        rffmpeg_ffmpeg_command.append(config["fallback_ffmpeg_command"])
        stdout = sys.stderr

    # Check for special flags that override the default stdout
    if any(item in config["special_flags"] for item in ffmpeg_args):
        stdout = sys.stdout

    # Append all the passed arguments directly
    for arg in ffmpeg_args:
        rffmpeg_ffmpeg_command.append("{}".format(arg))

    log.info("Local command: {}".format(" ".join(rffmpeg_ffmpeg_command)))

    with dbconn(config) as cur:
        cur.execute(
            "INSERT INTO processes (host_id, process_id, cmd) VALUES (?, ?, ?)",
            (0, config["current_pid"], cmd_name + " " + " ".join(ffmpeg_args)),
        )
        cur.execute(
            "INSERT INTO states (host_id, process_id, state) VALUES (?, ?, ?)",
            (0, config["current_pid"], "active"),
        )

    return run_command(rffmpeg_ffmpeg_command, stdin, stdout, stderr)


def run_remote_ffmpeg(config, target_hid, target_host, ffmpeg_args):
    """
    Run ffmpeg against the remote target_host.
    """
    rffmpeg_ssh_command = generate_ssh_command(config, target_host)
    rffmpeg_ffmpeg_command = list()

    # Add any pre commands
    for cmd in config["pre_commands"]:
        if cmd:
            rffmpeg_ffmpeg_command.append(cmd)

    # Prepare our default stdin/stderr
    stdin = sys.stdin
    stderr = sys.stderr

    if cmd_name == "ffprobe":
        # If we're in ffprobe mode use that command and sys.stdout as stdout
        rffmpeg_ffmpeg_command.append(config["ffprobe_command"])
        stdout = sys.stdout
    else:
        # Otherwise, we use stderr as stdout
        rffmpeg_ffmpeg_command.append(config["ffmpeg_command"])
        stdout = sys.stderr

    # Check for special flags that override the default stdout
    if any(item in config["special_flags"] for item in ffmpeg_args):
        stdout = sys.stdout

    # Append all the passed arguments with requoting of any problematic characters
    for arg in ffmpeg_args:
        # Match bad shell characters: * ' ( ) | [ ] or whitespace
        if search("[*'()|\[\]\s]", arg):
            rffmpeg_ffmpeg_command.append('"{}"'.format(arg))
        else:
            rffmpeg_ffmpeg_command.append("{}".format(arg))

    log.info(
        "Remote command: {}".format(
            " ".join(rffmpeg_ssh_command + rffmpeg_ffmpeg_command)
        )
    )

    with dbconn(config) as cur:
        cur.execute(
            "INSERT INTO processes (host_id, process_id, cmd) VALUES (?, ?, ?)",
            (target_hid, config["current_pid"], cmd_name + " " + " ".join(ffmpeg_args)),
        )
        cur.execute(
            "INSERT INTO states (host_id, process_id, state) VALUES (?, ?, ?)",
            (target_hid, config["current_pid"], "active"),
        )

    return run_command(
        rffmpeg_ssh_command + rffmpeg_ffmpeg_command, stdin, stdout, stderr
    )


def run_ffmpeg(config, ffmpeg_args):
    """
    Entrypoint for an ffmpeg/ffprobe aliased process.
    """
    signal.signal(signal.SIGTERM, cleanup)
    signal.signal(signal.SIGINT, cleanup)
    signal.signal(signal.SIGQUIT, cleanup)
    signal.signal(signal.SIGHUP, cleanup)

    if config["log_to_file"]:
        logging.basicConfig(
            filename=config["logfile"],
            level=logging.INFO,
            format="%(asctime)s - %(name)s[%(process)s] - %(levelname)s - %(message)s",
        )
    else:
        logging.basicConfig(
            level=logging.INFO,
            format="%(asctime)s - %(name)s[%(process)s] - %(levelname)s - %(message)s",
        )

    log.info("Starting rffmpeg with args: {}".format(" ".join(ffmpeg_args)))

    target_hid, target_hostname = get_target_host(config)

    if not target_hostname or target_hostname == "localhost":
        retcode = run_local_ffmpeg(config, ffmpeg_args)
    else:
        retcode = run_remote_ffmpeg(config, target_hid, target_hostname, ffmpeg_args)

    cleanup()
    if retcode == 0:
        log.info("Finished rffmpeg with return code {}".format(retcode))
    else:
        log.error("Finished rffmpeg with return code {}".format(retcode))
    exit(retcode)


def run_control(config):
    """
    Entrypoint for the Click CLI for managing the rffmpeg system.
    """
    CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"], max_content_width=120)

    @click.group(context_settings=CONTEXT_SETTINGS)
    def rffmpeg_click():
        """
        rffmpeg CLI interface
        """
        pass

    @click.command(name="init", short_help="Initialize the system.")
    @click.option(
        "-y",
        "--yes",
        "confirm_flag",
        is_flag=True,
        default=False,
        help="Confirm initialization.",
    )
    def rffmpeg_click_init(confirm_flag):
        """
        Initialize the rffmpeg system and database; this will erase all hosts and current state.

        This command should be run as "sudo" before any attempts to use rffmpeg.
        """
        if os.getuid() != 0:
            click.echo("Error: This command requires root privileges.")
            exit(1)

        if not confirm_flag:
            try:
                click.confirm(
                    "Are you sure you want to (re)initalize the database",
                    prompt_suffix="? ",
                    abort=True,
                )
            except Exception:
                fail("Aborting due to failed confirmation.")

        if not Path(config["state_dir"]).is_dir():
            try:
                os.makedirs(config["state_dir"])
            except OSError as e:
                fail(
                    "Failed to create state directory '{}': {}".format(
                        config["state_dir"], e
                    )
                )

        if Path(config["db_path"]).is_file():
            os.remove(config["db_path"])

        try:
            with dbconn(config) as cur:
                cur.execute(
                    """CREATE TABLE hosts (id INTEGER PRIMARY KEY, hostname TEXT NOT NULL, weight INTEGER DEFAULT 1)"""
                )
                cur.execute(
                    """CREATE TABLE processes (id INTEGER PRIMARY KEY, host_id INTEGER, process_id INTEGER, cmd TEXT)"""
                )
                cur.execute(
                    """CREATE TABLE states (id INTEGER PRIMARY KEY, host_id INTEGER, process_id INTEGER, state TEXT)"""
                )
        except Exception as e:
            fail("Failed to create database: {}".format(e))

        os.chown(
            config["state_dir"],
            getpwnam(config["dir_owner"]).pw_uid,
            getgrnam(config["dir_group"]).gr_gid,
        )
        os.chmod(config["state_dir"], 0o770)
        os.chown(
            config["db_path"],
            getpwnam(config["dir_owner"]).pw_uid,
            getgrnam(config["dir_group"]).gr_gid,
        )
        os.chmod(config["db_path"], 0o660)

    rffmpeg_click.add_command(rffmpeg_click_init)

    @click.command(name="status", short_help="Show hosts and status.")
    def rffmpeg_click_status():
        """
        Show the current status of all rffmpeg target hosts and active processes.
        """
        with dbconn(config) as cur:
            hosts = cur.execute("SELECT * FROM hosts").fetchall()
            processes = cur.execute("SELECT * FROM processes").fetchall()
            states = cur.execute("SELECT * FROM states").fetchall()

        # Determine if there are any fallback processes running
        fallback_processes = list()
        for process in processes:
            if process[1] == 0:
                fallback_processes.append(process)

        # Generate a mapping dictionary of hosts and processes
        host_mappings = dict()

        if len(fallback_processes) > 0:
            host_mappings[0] = {
                "hostname": "localhost (fallback)",
                "weight": 0,
                "current_state": "fallback",
                "commands": fallback_processes,
            }

        for host in hosts:
            hid, hostname, weight = host

            # Get the latest state
            with dbconn(config) as cur:
                current_state = cur.execute(
                    "SELECT * FROM states WHERE host_id = ? ORDER BY id DESC", (hid,)
                ).fetchone()

            if not current_state:
                current_state = "idle"
            else:
                current_state = current_state[3]

            # Create the mappings entry
            host_mappings[hid] = {
                "hostname": hostname,
                "weight": weight,
                "current_state": current_state,
                "commands": [proc for proc in processes if proc[1] == hid],
            }

        hostname_length = 9
        hid_length = 3
        weight_length = 7
        state_length = 6
        for hid, host in host_mappings.items():
            if len(host["hostname"]) + 1 > hostname_length:
                hostname_length = len(host["hostname"]) + 1
            if len(str(hid)) + 1 > hid_length:
                hid_length = len(str(hid)) + 1
            if len(str(host["weight"])) + 1 > weight_length:
                weight_length = len(str(host["weight"])) + 1
            if len(host["current_state"]) + 1 > state_length:
                state_length = len(host["current_state"]) + 1

        output = list()
        output.append(
            "{bold}{hostname: <{hostname_length}} {hid: <{hid_length}} {weight: <{weight_length}} {state: <{state_length}} {commands}{end_bold}".format(
                bold="\033[1m",
                end_bold="\033[0m",
                hostname="Hostname",
                hostname_length=hostname_length,
                hid="ID",
                hid_length=hid_length,
                weight="Weight",
                weight_length=weight_length,
                state="State",
                state_length=state_length,
                commands="Active Commands",
            )
        )

        for hid, host in host_mappings.items():
            if len(host["commands"]) < 1:
                first_command = "N/A"
            else:
                first_command = "PID {}: {}".format(
                    host["commands"][0][2], host["commands"][0][3]
                )

            host_entry = list()
            host_entry.append(
                "{hostname: <{hostname_length}} {hid: <{hid_length}} {weight: <{weight_length}} {state: <{state_length}} {commands}".format(
                    hostname=host["hostname"],
                    hostname_length=hostname_length,
                    hid=hid,
                    hid_length=hid_length,
                    weight=host["weight"],
                    weight_length=weight_length,
                    state=host["current_state"],
                    state_length=state_length,
                    commands=first_command,
                )
            )

            for idx, command in enumerate(host["commands"]):
                if idx == 0:
                    continue
                host_entry.append(
                    "{hostname: <{hostname_length}} {hid: <{hid_length}} {weight: <{weight_length}} {state: <{state_length}} {commands}".format(
                        hostname="",
                        hostname_length=hostname_length,
                        hid="",
                        hid_length=hid_length,
                        weight="",
                        weight_length=weight_length,
                        state="",
                        state_length=state_length,
                        commands="PID {}: {}".format(command[2], command[3]),
                    )
                )

            output.append("\n".join(host_entry))

        click.echo("\n".join(output))

    rffmpeg_click.add_command(rffmpeg_click_status)

    @click.command(name="add", short_help="Add a host.")
    @click.option(
        "-w",
        "--weight",
        "weight",
        required=False,
        default=1,
        help="The weight of the host.",
    )
    @click.argument("host")
    def rffmpeg_click_add(weight, host):
        """
        Add a new host with IP or hostname HOST to the database.
        """
        with dbconn(config) as cur:
            cur.execute(
                """INSERT INTO hosts (hostname, weight) VALUES (?, ?)""", (host, weight)
            )

    rffmpeg_click.add_command(rffmpeg_click_add)

    @click.command(name="remove", short_help="Remove a host.")
    @click.argument("host")
    def rffmpeg_click_remove(host):
        """
        Remove a host with internal ID or IP or hostname HOST from the database.
        """
        try:
            host = int(host)
            field = "id"
        except ValueError:
            field = "hostname"

        with dbconn(config) as cur:
            entry = cur.execute(
                "SELECT * FROM hosts WHERE {} = ?".format(field), (host,)
            ).fetchall()
            if len(entry) < 1:
                fail("No hosts found to delete!")

            click.echo("Deleting {} host(s):".format(len(entry)))
            for h in entry:
                click.echo("\tID: {}\tHostname: {}".format(h[0], h[1]))
                cur.execute("""DELETE FROM hosts WHERE id = ?""", (h[0],))

    rffmpeg_click.add_command(rffmpeg_click_remove)

    @click.command(name="log", short_help="View the rffmpeg log.")
    @click.option(
        "-f",
        "--follow",
        "follow_flag",
        is_flag=True,
        default=False,
        help="Follow new log entries instead of seeing current log entries.",
    )
    def rffmpeg_click_log(follow_flag):
        """
        View the rffmpeg log file.
        """
        if follow_flag:
            with open(config["logfile"]) as file_:
                # Go to the end of file
                file_.seek(0, 2)
                while True:
                    curr_position = file_.tell()
                    line = file_.readline()
                    if not line:
                        file_.seek(curr_position)
                        sleep(0.1)
                    else:
                        click.echo(line, nl=False)
        else:
            with open(config["logfile"], "r") as logfh:
                click.echo_via_pager(logfh.readlines())

    rffmpeg_click.add_command(rffmpeg_click_log)

    return rffmpeg_click(obj={})


# Entrypoint
if __name__ == "__main__":
    all_args = sys.argv
    cmd_name = all_args[0]

    # Load the config
    config = load_config()

    if "rffmpeg" in cmd_name:
        run_control(config)
    else:
        if not Path(config["db_path"]).is_file():
            fail(
                "Failed to find database '{}' - did you forget to run 'rffmpeg init'?".format(
                    config["db_path"]
                )
            )

        ffmpeg_args = all_args[1:]
        run_ffmpeg(config, ffmpeg_args)