ffmpeg-transcode

#!/usr/bin/env python3

import re
import yaml
import os
import sys
import subprocess
import platform
import signal
from shlex import join
from shlex import split
from datetime import datetime

# ------------------------------------------------------------------------------
# Resolve absolute path (if given path is relative)
def resolve_absolute_path(filename):
  if filename.startswith("/"):
    # Absolute path
    return filename
  else:
    # Relative path
    return os.path.dirname(os.path.realpath(__file__))+"/"+filename

# ------------------------------------------------------------------------------
# Load configuration file
def load_configuration():
  #print("Loading configuration...")
  config = None
  configFile = resolve_absolute_path("ffmpeg-transcode.yaml")
  configFallback = "ffmpeg-transcode-example.yaml"
  if os.path.isdir('/config'):
    # Use /config if existing (for docker support)
    configFile = "/config/ffmpeg-transcode.yaml"
    configFallback = "ffmpeg-transcode-docker.yaml"
  if not os.path.isfile(configFile):
    # Download default config file via wget
    subprocess.run(["wget", "-q", "https://github.com/ForsakenNGS/raspi-plex-transcode/raw/main/"+configFallback, "-O", configFile]);
  # Open and parse configuration file
  with open(configFile, "r") as stream:
    try:
      config = yaml.safe_load(stream)
    except yaml.YAMLError as e:
      print("Error loading config: "+configFile)
      print(e)
  return config

def parse_argument_value(value):
  if value is None:
    return None
  else:
    return " ".join(value)

# ------------------------------------------------------------------------------
# Parse default arguments
def parse_arguments():
  #print("Parsing commandline arguments...")
  arguments = {
    "input": {},
    "output": {}
  }
  argSection = "input"
  argName = None
  argValue = None
  for arg in sys.argv[1:]:
    if arg.startswith("-"):
      if argName is not None:
        arguments[argSection][argName] = parse_argument_value(argValue)
        if argName == "-i":
          argSection = "output"
      argName = arg
      argValue = None
    else:
      if argValue is None:
        argValue = []
      argValue.append(arg)
  if argName is not None:
    arguments[argSection][argName] = parse_argument_value(argValue)
  return arguments

def profile_add(profilesByGroup, profile, profileGroup, priority):
  profilesByGroup[profileGroup] = profile;
  profilesByGroup[profileGroup]["priority"] = priority;
  return profilesByGroup

# ------------------------------------------------------------------------------
# Check profile condition and add target profile if applicable
def profile_condition_arg(condition, arguments, config, profilesByGroup):
  # Get target profile settings
  profileTargetName = condition["profile"]
  profileTarget = config["profiles"][profileTargetName]
  profileGroup = "general"
  if "group" in profileTarget:
    profileGroup = profileTarget["group"]
  # Get current profile
  currentProfile = None
  currentPriority = -1
  if profileGroup in profilesByGroup:
    currentProfile = profilesByGroup[profileGroup]
    if "priority" in currentProfile:
      currentPriority = currentProfile["priority"]
  # Check priority
  conditionPriority = 0
  if "priority" in condition:
    conditionPriority = condition["priority"]
  if conditionPriority <= currentPriority:
    return profilesByGroup
  # Get argument value
  if not condition["argName"] in arguments[condition["argSection"]]:
    if condition["type"] == "missing":
      # Matches if argument is missing
      return profile_add(profilesByGroup, profileTarget, profileGroup, conditionPriority)
    else:
      # Condition type requires argument to be present
      return profilesByGroup
  value = arguments[condition["argSection"]][condition["argName"]]
  # Check condition
  if condition["type"] == "present":
    # Matches as soon as the argument is present
    return profile_add(profilesByGroup, profileTarget, profileGroup, conditionPriority)
  elif condition["type"] == "exact":
    # Exact value match
    if condition["value"] == value:
      return profile_add(profilesByGroup, profileTarget, profileGroup, conditionPriority)
  elif condition["type"] == "regex":
    # Regular expression match
    regexFlags = 0
    if ("ignorecase" in condition) and condition["ignorecase"]:
      regexFlags += re.IGNORECASE
    regexPattern = re.compile(condition["value"], regexFlags)
    if regexPattern.match(value):
      return profile_add(profilesByGroup, profileTarget, profileGroup, conditionPriority)
  return profilesByGroup

# ------------------------------------------------------------------------------
# Select profiles to be used based on config and default arguments
def profile_select(config, arguments):
  #print("Evaluating target profile...")
  profilesByGroup = {}
  if not "profile_select" in config:
    return None
  # Apply default profile
  if "default" in config["profile_select"]:
    profileTarget = config["profiles"][ config["profile_select"]["default"] ]
    profileGroup = "general"
    if "group" in profileTarget:
      profileGroup = profileTarget["group"]
    profile_add(profilesByGroup, profileTarget, profileGroup, 0)
  # Process profile conditions
  # -> By Argument
  if "by_argument" in config["profile_select"]:
    for condition in config["profile_select"]["by_argument"]:
      profilesByGroup = profile_condition_arg(condition, arguments, config, profilesByGroup)
  # -> By Codec
  if "by_codec" in config["profile_select"]:
    probeBin = resolve_absolute_path(config["ffprobe"])
    probeArgs = [ probeBin, "-v", "error", "-show_entries", "stream=codec_name", "-of", "default=noprint_wrappers=1:nokey=1", arguments["input"]["-i"] ]
    probeResult = subprocess.run(probeArgs, universal_newlines = True, stdout = subprocess.PIPE)
    probeCodecs = probeResult.stdout.splitlines()
    for probeIndex, probeCodec in enumerate(probeCodecs):
      for condition in config["profile_select"]["by_codec"]:
        if condition["source"] == probeCodec:
          conditionPriority = 0
          if "priority" in condition:
            conditionPriority = condition["priority"]
          profileCodecGroup = "codec"+str(probeIndex)
          profileCodec = {
            "group": profileCodecGroup,
            "input": {},
            "output": {
              "-codec:"+str(probeIndex): condition["target"]
            }
          }
          if "input" in condition:
            for inputArg, inputVal in condition["input"].items():
              profileCodec["input"][inputArg.replace("$stream$", str(probeIndex))] = inputVal
          if "output" in condition:
            for inputArg, inputVal in condition["output"].items():
              profileCodec["output"][inputArg.replace("$stream$", str(probeIndex))] = inputVal
          profile_add(profilesByGroup, profileCodec, profileCodecGroup, conditionPriority)
  # List and sort all selected profiles
  profileList = []
  for profileGroup in profilesByGroup:
    profileList.append(profilesByGroup[profileGroup])
  def profileSort(profile):
    return profile["priority"]
  profileList.sort(key=profileSort)
  return profileList

# ------------------------------------------------------------------------------
# Modify parameters based on the profiles given
def profile_parameters(profiles):
  argumentRepl = {
    "input": {},
    "output": {}
  }
  for profile in profiles:
    # Apply profile overrides
    if profile["input"] is not None:
      for argName in profile["input"]:
        argumentRepl["input"][argName] = profile["input"][argName]
    if profile["output"] is not None:
      for argName in profile["output"]:
        argumentRepl["output"][argName] = profile["output"][argName]
  # Build parameter list
  parameters = []
  argSection = "input"
  argName = None
  argNewAdded = False
  for arg in sys.argv[1:]:
    # Apply replacement if present
    replacement = None
    if argName is not None:
      if argName in argumentRepl[argSection]:
        replacement = argumentRepl[argSection][argName]
        del argumentRepl[argSection][argName]
    if replacement is not None:
      if replacement == "**REMOVE**":
        del parameters[-1]
        argName = None
        if not arg.startswith("-"):
          continue
      elif not arg.startswith("-"):
        arg = replacement
    # Check for end of input/output arguments
    if not argNewAdded and ((arg == "-i") or (arg == "-init_seg_name") or (arg == "-segment_header_filename")):
      # Add parameters that were not present in the inital call from plex
      for argNameNew in argumentRepl[argSection]:
        argValueNew = argumentRepl[argSection][argNameNew]
        if argValueNew == "**REMOVE**":
          continue
        if not argNameNew in sys.argv:
          parameters.append(argNameNew)
          if argValueNew is not None:
            parameters.append(argValueNew)
      argNewAdded = True
    # Append to parameter list
    parameters.append(arg)
    # Update state
    if arg.startswith("-"):
      argName = arg
    else:
      if argName == "-i":
        argSection = "output"
        argNewAdded = False
      argName = None
  return parameters

# ------------------------------------------------------------------------------
# Start process for port-fowarding from transcode proxy to the local plex instance
def proxy_tunnel(proxyConfig):
  proxyHost = proxyConfig["host"]
  proxyUser = proxyConfig["user"]
  tunnelArgs = [ "ssh", "-N", "-R", "32400:127.0.0.1:32400", "-o", "ExitOnForwardFailure=yes", proxyUser+"@"+proxyHost ]
  tunnelProc = subprocess.Popen(tunnelArgs, start_new_session=True)
  try:
    tunnelProc.wait(timeout=2)
    return False    # Failed to create tunnel, assume already open
  except subprocess.TimeoutExpired:
    return True     # New tunnel established

# ------------------------------------------------------------------------------
# Transcode stream on a proxy transcode-server (or return False to transcode locally)
def proxy_transcode(config):
  if not "proxys" in config:
    return False
  for proxyConfig in config["proxys"]:
    # Local configuration
    plexHostname = platform.node()
    # Parse configuration
    proxyType = "ssh"
    if "type" in proxyConfig:
      proxyType = proxyConfig["type"]
    proxyHost = proxyConfig["host"]
    proxyUser = proxyConfig["user"]
    proxyCmdsPre = []
    proxyCmdsPost = []
    if "cmds-pre" in proxyConfig:
      proxyCmdsPre = proxyConfig["cmds-pre"]
    if "cmds-post" in proxyConfig:
      proxyCmdsPost = proxyConfig["cmds-post"]
    proxyCmdsPreFailed = False
    # Execute pre-commands
    for cmdPre in proxyCmdsPre:
      cmdArgs = split(cmdPre)
      cmdResult = subprocess.run(cmdArgs)
      if cmdResult.returncode > 0:
        proxyCmdsPreFailed = True
        break
    if proxyCmdsPreFailed:
      continue
    # Execute stream
    if proxyType == "ssh":
      proxy_tunnel(proxyConfig)
      cmdArgs = [ "ssh" ]
      if "ssh-options" in proxyConfig:
        for cmdSshArg in proxyConfig["ssh-options"]:
          cmdArgs.append(cmdSshArg)
      cmdArgs.append(proxyUser+"@"+proxyHost)
      cmdFfmpegArgs = sys.argv[1:]
      cmdFfmpegArgs.insert(0, proxyConfig["ffmpeg"])
      cmdArgs.append(join(cmdFfmpegArgs))

      if ("log" in config) and ("debug" in config):
        with open(config["log"], "a") as logfile:
          logfile.write("Starting proxy-transcode at "+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+"\n")
          logfile.write("Commandline: "+join(cmdArgs)+"\n\n\n")

      cmdResult = subprocess.run(cmdArgs)
      if cmdResult.returncode > 0:
        if "log" in config:
          with open(config["log"], "a") as logfile:
            logfile.write("Failed proxy-transcode at "+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+"\n")
            logfile.write("Commandline: "+join(cmdArgs)+"\n")
            logfile.write("Return code: "+str(cmdResult.returncode)+"\n")
            if cmdResult.stderr is not None:
              logfile.write("STDERR Output: ")
              logfile.write(cmdResult.stderr)
            logfile.write("\n\n\n")
        continue
      else:
        # Execute post-commands
        for cmdPost in proxyCmdsPost:
          cmdArgs = split(cmdPost)
          cmdResult = subprocess.run(cmdArgs)
          if cmdResult.returncode > 0:
            break
        return True
  # No proxy transcode executed (successfully)
  return False

# ------------------------------------------------------------------------------
# Entry point

# Load configuration
config = load_configuration()
if config is None:
  print("Failed to load configuration! Exiting.")
  sys.exit(1)

# Parse arguments
arguments = parse_arguments()

# Try to pass transcode to proxy
if not proxy_transcode(config):
  # No proxy configured / available
  # Select matching profiles
  profiles = profile_select(config, arguments)
  # Get modified parameters
  parameters = profile_parameters(profiles)
  # Prepare transcode call
  executable = resolve_absolute_path(config["ffmpeg"])
  parameters.insert(0, executable)
  # Log transcode call (if debug enabled)
  if ("log" in config) and ("debug" in config):
    with open(config["log"], "a") as logfile:
      logfile.write("Starting transcode at "+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+"\n")
      logfile.write("Commandline: "+join(parameters)+"\n\n\n")
  # Start transcode process
  result = subprocess.run(parameters, stderr = subprocess.PIPE)
  # Log transcode fail
  if result.returncode > 0:
    if "log" in config:
      with open(config["log"], "a") as logfile:
        logfile.write("Failed transcode at "+datetime.now().strftime("%Y-%m-%d %H:%M:%S")+"\n")
        logfile.write("Commandline: "+join(parameters)+"\n")
        logfile.write("Return code: "+str(result.returncode)+"\n")
        logfile.write("STDERR Output: ")
        if result.stderr is not None:
          logfile.write("STDERR Output: ")
          logfile.write(result.stderr)
        logfile.write("\n\n\n")
  # Exit with same return code as the transcode process
  sys.exit(result.returncode)
else:
  sys.exit(0)