mpegts

#!/usr/bin/env python
"""
mpegts v1.4 - Copyright (C) 2016 Ciriaco Garcia de Celis
Utility to handle multiplexed MPEG-TS files captured from raw DVB sources
"""
# 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 <http://www.gnu.org/licenses/>.

import os
import stat
import sys
import subprocess
import re
import textwrap
from datetime import timedelta
from collections import OrderedDict, defaultdict
from itertools import chain

inputMode = os.fstat(sys.stdin.fileno()).st_mode
fromStdin = stat.S_ISFIFO(inputMode) or stat.S_ISREG(inputMode)
varg = 1 if fromStdin else 2

if len(sys.argv) < varg:
    print(textwrap.indent(__doc__, '  ') + '\n'
          '    ' + os.path.basename(sys.argv[0]) + ' input.mts [channel [operation ...]]\n\n'
          '      channel (omit to list all channels):\n'
          '          "channel_name" or channel_ID\n\n'
          '      operation (omit to show the streams in the selected channel):\n'
          '          play                 [t0 [t1 [t2]]  (reproduce the channel)\n'
          '          save     output.ts   [t0 [t1 [t2]]  (extract channel; omits teletext)\n'
          '          save_ns  output.ts   [t0 [t1 [t2]]  (also omits the subtitles)\n'
          '          compress output.mkv  [t0 [t1 [t2]]  (HQ vp9 encoding; with subtitles)\n'
          '          compress output.webm [t0 [t1 [t2]]  (HQ vp9 encoding; subtitles lost)\n'
          '          compress output      [t0 [t1 [t2]]  (use .mkv if subtitles, or .webm)\n\n'
          '            t0: stream absolute start time in seconds (may be set to 0)\n'
          '            t1: crop absolute start time (t1-t0 = relative start): 0 by default\n'
          '            t2: crop absolute end time (t2-t1 = duration): stream end by default\n\n'
          '      Note: Input and/or output files must be omitted when redirection is used\n')
    sys.exit(1)

inFile = '-' if fromStdin else sys.argv[1]
probesize = '10M'

proc = subprocess.Popen(['ffprobe',
                         '-v', 'quiet',
                         '-probesize', probesize,
                         '-of', 'compact',
                         '-show_programs',  # also reports streams (without subtitles lang if -show_streams isn't set)
                         '-show_streams',  # this duplicates streams listing when -show_programs is set
                         '-show_entries', 'format',  # for duration (when not available per stream)
                         '-' if fromStdin else 'file:' + inFile],
                        stdout=subprocess.PIPE)

channels = {}
channelStreams = defaultdict(list)
streams = {}
lastChannelName = None
nb_streams = 0
fileDuration = None

while True:
    input = proc.stdout.readline()
    if not input:
        break
    line = re.sub('\\x05|\\n', '', input.decode("utf-8"))
    if len(line) < 1:
        nb_streams = 0  # ffprobe bug in programs lookup with some files
        continue
    data = dict([(x.split('=') + [None])[:2] for x in line.split('|')])
    if 'format' in data:
        if 'duration' in data:
            fileDuration = data['duration'] if data['duration'] != 'N/A' else None
    elif 'program_id' in data:
        if nb_streams == 0:
            nb_streams = int(data['nb_streams'])
        if 'tag:service_name' in data:
            lastChannelName = data['tag:service_name'].lower()
            channels[lastChannelName] = data
    elif nb_streams > 0:
        nb_streams -= 1
        channelStreams[lastChannelName].append(data['index'])
    elif 'index' in data:
        streams[data['index']] = data

proc.wait()
if proc.returncode != 0:
    print('error reading', inFile)
    sys.exit(proc.returncode)

channelName = sys.argv[varg] if len(sys.argv) > varg and len(channelStreams) > 0 else None

if len(channelStreams) == 0:
    varg = varg - 1

if channels and not channelName:
    longestChannel = max(len(x) for x in channels)
    sortedChannels = OrderedDict(sorted(channels.items(), key=lambda i: i[0].casefold()))
    for channel in sortedChannels.values():
        channelName = channel['tag:service_name']
        nameTellsProvider = re.match(channel['tag:service_provider'], channelName, re.IGNORECASE)
        print('ID {0:>5} ({1:>2} streams) {2:{3}}{4}'.format(
              channel['program_id'],
              channel['nb_streams'],
              channelName, longestChannel,
              '' if nameTellsProvider else ' (%s)' % channel['tag:service_provider']
              ))
    sys.exit(0)

if channelName and channelName.isnumeric():
    channelName = dict(((prop['program_id'], prop['tag:service_name'])
                       for name, prop in channels.items())).get(channelName, channelName)

videoPids = []
audioPids = []
subtitlePids = []

showStreams = not len(sys.argv) > varg + 1

if channelName and channelName.lower() not in channelStreams:
    print('channel "%s" not found' % channelName)
    sys.exit(1)
else:
    if channelName and showStreams:
        print(channelName)
    for streamId in (channelStreams[channelName.lower()] if channelName else streams.keys()):
        stream = streams[streamId]
        if 'codec_type' not in stream:
            continue
        isUnknown = stream['codec_type'] == 'unknown'
        isVideo = stream['codec_type'] == 'video'
        isAudio = stream['codec_type'] == 'audio'
        isImpaired = 'disposition:visual_impaired' in stream and stream['disposition:visual_impaired'] == '1'
        if isUnknown or (isAudio and isImpaired):
            continue
        pid = None if 'id' not in stream or stream['id'] == 'N/A' else stream['id']
        if showStreams:
            duration = stream['duration'] if 'duration' in stream and stream['duration'] != 'N/A' else None
            if not duration and (isVideo or isAudio):
                duration = fileDuration
            print("{0} {1}={2} {3:4} [{4}]".format(
                stream['codec_type'],
                'PID' if pid else 'IDX',
                pid if pid else stream['index'],
                stream['codec_name'],
                "{:0>8}".format(str(timedelta(seconds=int(float(duration))))) if duration else 'N/A'),
                end=' ')
        if isVideo:
            if pid:
                videoPids.append(pid)
            if showStreams:
                print(stream['width'] + 'x' + stream['height'],
                      '' if 'display_aspect_ratio' not in stream else stream['display_aspect_ratio'])
        elif isAudio:
            if pid:
                audioPids.append(pid)
            if showStreams:
                print('' if 'tag:language' not in stream else '(%s)' % stream['tag:language'],
                      '' if 'channel_layout' not in stream else stream['channel_layout'],
                      '' if 'channels' not in stream else stream['channels'],
                      '' if 'bit_rate' not in stream else stream['bit_rate'])
        elif stream['codec_type'] == 'subtitle':
            if pid and 'dvbsub' in stream['codec_name'].replace('_', ''):
                subtitlePids.append(pid)
            if showStreams:
                print('' if 'tag:language' not in stream else '(%s)' % stream['tag:language'])

if showStreams:
    sys.exit(0)

operation = sys.argv[varg + 1].lower()
isPlay = operation == 'play'
isSave = operation == 'save' or operation == 'save_ns'
isCompress = operation == 'compress'

if operation == 'save_ns':
    subtitlePids = []


def seek(argc, opt):
    return [] if len(sys.argv) <= argc else \
           [opt, str(float(sys.argv[argc]) - float(sys.argv[argc - 1]))]

if isPlay:
    cmd = ['ffplay']
    if len(videoPids) > 0 or len(audioPids) > 0:
        cmd.extend(['-vst', 'i:' + videoPids[0]] if len(videoPids) > 0 else ['-vn', '-nodisp'])
        cmd.extend(['-ast', 'i:' + audioPids[0]] if len(audioPids) > 0 else ['-an'])
    cmd.append('-sn')
    cmd.extend(seek(varg + 3, '-ss'))
    cmd.extend(seek(varg + 4, '-t'))
    cmd.append('-' if fromStdin else 'file:' + inFile)
elif (isSave or isCompress):
    toStdout = len(sys.argv) < varg + 3
    if toStdout:
        outputMode = os.fstat(sys.stdout.fileno()).st_mode
        if not (stat.S_ISFIFO(outputMode) or stat.S_ISREG(outputMode)):  # if output not redirected
            print('missing output file argument')
            sys.exit(1)
        output, ext = ('-', '')
    else:
        output, ext = os.path.splitext(sys.argv[varg + 2])
        ext = ext.lower()
        if isCompress and ext != '.mkv' and ext != '.webm':
            ext = '.mkv' if len(subtitlePids) > 0 else '.webm'
        varg += 1
    cmd = ['ffmpeg']
    cmd.extend(['-probesize', probesize])
    cmd.extend(['-i', '-' if fromStdin else 'file:' + inFile])
    cmd.extend(seek(varg + 3, '-ss'))
    cmd.extend(seek(varg + 4, '-t'))
    for pid in chain(videoPids, audioPids, subtitlePids):
        cmd.append('-map')
        cmd.append('i:' + pid)
    if isSave:
        cmd.extend(['-c', 'copy'])
    else:
        cmd.extend(['-threads', str(2 if sys.version_info < (3, 4) else os.cpu_count())])
        cmd.extend(['-c:v', 'libvpx-vp9',  # perceptually betters h265 at high compression
                    '-crf', '32', '-b:v', '2.5M',  # very high (but constrained) quality
                    '-g', '300'])  # proper seek in some players
        if ext == '.webm':
            cmd.extend(['-c:a', 'libopus', '-b:a', '96k'])  # webm requirement and much better than mp3
            cmd.extend(['-sn'])
        else:
            cmd.extend(['-c:a', 'copy'])
            cmd.extend(['-c:s', 'copy'])
        os.nice(19)
    if toStdout:
        cmd.extend(['-f', 'mpegts' if isSave else 'matroska', '-'])
    else:
        if fromStdin:
            cmd.append('-y')
        cmd.append('file:' + output + ext)
else:
    print('unknown operation "%s"' % operation)
    sys.exit(1)

retcode = subprocess.call(cmd)

if retcode != 0:
    sys.exit(retcode)
else:
    if (isSave or isCompress) and not (fromStdin or toStdout):
        finf = os.stat(inFile)
        os.utime(output + ext, (finf.st_atime, finf.st_mtime))
    sys.exit(0)