HDHRUtil-DVR-fileMetadata

#!/usr/bin/env python3

"""

HDHRUtil-DVR-fileMetadata

HDHRUtil-DVR-fileMetadata is a utility for manipulating
the HDHR DVR file metadata

Copyright (c) 2018 by Gary Buhrmaster <gary.buhrmaster@gmail.com>

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.



===== NOTE ===== NOTE ===== NOTE ===== NOTE ===== NOTE =====

This utility currently uses a reverse engineered
schema definition to the file metadata.  As it is
not documented as a stable ABI, it could change at
any time.

===== NOTE ===== NOTE ===== NOTE ===== NOTE ===== NOTE =====


"""


import sys
import json
import argparse
import re
import datetime
import dateutil
import dateutil.parser
import dateutil.tz


def episodeNumNormalize(episodeNum):
    m0 = re.match(r'^S(\d+)E(\d+)$', episodeNum, re.IGNORECASE)
    if m0:
        return 'S{0:02}E{1:02}'.format(int(m0.group(1)), int(m0.group(2)))
    else:
        raise TypeError('Invalid episode number: {0}'.format(episodeNum))

def episodeNumCheck(episodeNum):
    try:
        return episodeNumNormalize(episodeNum)
    except TypeError:
        raise argparse.ArgumentTypeError('{} is not a valid episode number (SnnEnn)'.format(episodeNum))

def programIDCheck(programID):
    if (len(programID) == 14) and (programID[0:2] in ['EP', 'MV', 'SH', 'SP']) and (re.match(r'^[0-9]+$', programID[2:])):
        return programID
    else:
        raise argparse.ArgumentTypeError('{} is not a valid ProgramID (14 characters, starts with EP, MV, SH, SP'.format(programID))

def datetimeCheck(dt):
    # Accept a epoch integer (presume the person knows for what they offer)
    if re.match(r'^(\d+)$', dt):
        return int(dt)
    try:
        d = dateutil.parser.parse(dt)
        if d.tzinfo is None:
            d = d.replace(tzinfo=dateutil.tz.tzlocal())
        return int((d - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)).total_seconds())
    except:
        raise argparse.ArgumentTypeError('{} is not a valid datetime'.format(dt))

def airdateCheck(dt):
    # Accept a epoch integer (presume the person knows for what they offer)
    if re.match(r'^(\d+)$', dt):
        return int(dt)
    try:
        d = dateutil.parser.parse(dt)
        if d.tzinfo is None:
            d = d.replace(tzinfo=dateutil.tz.tzlocal())
        dz = d.astimezone(tz=datetime.timezone.utc)
        dm = datetime.datetime(dz.year, dz.month, dz.day, tzinfo=datetime.timezone.utc)
        return int((dm - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)).total_seconds())
    except:
        raise argparse.ArgumentTypeError('{} is not a valid datetime'.format(dt))

def channelNormalize(channel):
    m0 = re.match(r'^(\d+)$', channel)
    m1 = re.match(r'^(\d+)\.(\d+)$', channel)
    m2 = re.match(r'^(\d+)_(\d+)$', channel)
    m3 = re.match(r'^(\d+)-(\d+)$', channel)
    if m0:
        return '{}'.format(int(m0.group(1)))
    elif m1:
        return '{}.{}'.format(int(m1.group(1)), int(m1.group(2)))
    elif m2:
        return '{}.{}'.format(int(m2.group(1)), int(m2.group(2)))
    elif m3:
        return '{}.{}'.format(int(m3.group(1)), int(m3.group(2)))
    raise TypeError('Invalid channel: {}'.format(channel))

def channelCheck(channel):
    try:
        return channelNormalize(channel)
    except Exception:
        raise argparse.ArgumentTypeError('{} is not a valid channel'.format(channel))

def transportStreamToMetadata(ts):
    # Obtain metadata from transport stream
    metadata = bytearray()
    if len(ts) != 12032:
        raise ValueError('transport stream too short for metadata')
    for pkt in range(0, 64, 1):
        if ts[pkt * 188 + 0] != 0x47:
            raise ValueError('transport stream does not contain proper sync bytes')
        if pkt == 0:
            if ts[pkt * 188 + 1] != 0x5f:
                raise ValueError('transport stream does not contain proper start payload and PID value bytes')
        else:
            if ts[pkt * 188 + 1] != 0x1f:
                raise ValueError('transport stream does not contain proper payload and PID value bytes')
        if ts[pkt * 188 + 2] != 0xfa:
            raise ValueError('transport stream does not contain proper PID value')
        if ts[pkt * 188 + 3] != 0x10 + pkt % 16:
            raise ValueError('transport stream does not contain proper adaption and sequence bytes')
        for c in range(4, 188, 1):
            if ts[pkt * 188 + c] == 0xff:
                break
            metadata.append(ts[pkt * 188 + c])
    if len(metadata) <= 0:
        raise ValueError('transport stream does not contain any metadata')
    try:
        return json.loads(metadata)
    except:
        raise ValueError('transport stream metadata is not decodable as json')

def metadataToTransportStream(metadata):
    # Create transport stream with metadata
    TS = bytearray([0xff]*12032)
    for pkt in range(0, 64, 1):
        TS[pkt * 188 + 0] = 0x47               # sync byte
        if pkt == 0:
            TS[pkt * 188 + 1] = 0x5f           # start + PID (0x1ffa)
        else:
            TS[pkt * 188 + 1] = 0x1f           # PID (0x1ffa)
        TS[pkt * 188 + 2] = 0xfa
        TS[pkt * 188 + 3] = 0x10 + pkt % 16    # Scrambling (none) + Adaption (payload only) + sequence
    if type(metadata) != dict:
        raise TypeError('metadata is not a valid dict')
    try:
        meta = bytearray(json.dumps(metadata, separators=(',', ':')), 'utf-8')
    except:
        raise ValueError('metadata can not be converted to json')
    if len(meta) > 11776:
        raise ValueError('metadata json string exceeds allowable length')
    for pkt in range(0, 64, 1):
        l = min(len(meta), 184)
        offset = pkt * 188 + 4
        for c in range(0, l, 1):
            TS[offset + c] = meta[c]
        meta = meta[l:]
        if len(meta) <= 0:
            break
    return TS


if __name__ == '__main__':

    # Parse our args
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', action='store', dest='file', required=True,
                        help='The file whose metadata will be manipulated')
    actiongroup = parser.add_mutually_exclusive_group(required=False)
    actiongroup.add_argument('--update', action='store_true', dest='update', default=False,
                             help='Update the file metadata')
    actiongroup.add_argument('--set', action='store_true', dest='set', default=False,
                             help='Set the file metadata')
    actiongroup.add_argument('--delete', nargs='+', dest='delete',
                             choices=['Category', 'ChannelAffiliate', 'ChannelImageURL',
                                      'ChannelName', 'EndTime', 'EpisodeNumber', 'EpisodeTitle',
                                      'FirstAiring', 'ImageURL', 'OriginalAirdate',
                                      'ProgramID', 'RecordEndTime', 'RecordStartTime',
                                      'RecordSuccess', 'Resume', 'SeriesID', 'StartTime',
                                      'Synopsis', 'Team1', 'Team2', 'Title'],
                             help='Delete the specified keys in the file metadata')
    parser.add_argument('--Category', action='store', dest='Category',
                        choices=['series', 'movie', 'news', 'sports', 'special', 'other'],
                        help='The program category (one of series, movie, news, sports, special, or other)')
    parser.add_argument('--ChannelAffiliate', action='store', dest='ChannelAffiliate',
                        help='The program channel affiliate')
    parser.add_argument('--ChannelImageURL', action='store', dest='ChannelImageURL',
                        help='The program channel image url')
    parser.add_argument('--ChannelName', action='store', dest='ChannelName',
                        help='The program channel name')
    parser.add_argument('--ChannelNumber', action='store', type=channelCheck, dest='ChannelNumber',
                        help='The program channel number')
    parser.add_argument('--EndTime', action='store', type=datetimeCheck, dest='EndTime',
                        help='The program end time')
    parser.add_argument('--EpisodeNumber', action='store', type=episodeNumCheck, dest='EpisodeNumber',
                        help='The program episode number (SnnEnn)')
    parser.add_argument('--EpisodeTitle', action='store', dest='EpisodeTitle',
                        help='The program episode title')
    parser.add_argument('--FirstAiring', action='store_const', const=1, dest='FirstAiring',
                        help='Mark the recording as a first airing')
    parser.add_argument('--no-FirstAiring', action='store_const', const=0, dest='FirstAiring',
                        help='Mark the recording as not a first airing')
    parser.add_argument('--ImageURL', action='store', dest='ImageURL',
                        help='The program image')
    parser.add_argument('--OriginalAirdate', action='store', type=airdateCheck, dest='OriginalAirdate',
                        help='The program original air date')
    parser.add_argument('--ProgramID', action='store', dest='ProgramID',
                        help='The program id')
    parser.add_argument('--RecordEndTime', action='store', type=datetimeCheck, dest='RecordEndTime',
                        help='The program recording end time')
    parser.add_argument('--RecordStartTime', action='store', type=datetimeCheck, dest='RecordStartTime',
                        help='The program recording start time')
    parser.add_argument('--RecordSuccess', action='store_const', const=1, dest='RecordSuccess',
                        help='Mark the recording as successful')
    parser.add_argument('--no-RecordSuccess', action='store_const', const=0, dest='RecordSuccess',
                        help='Mark the recording as not successful')
    parser.add_argument('--Resume', action='store', dest='Resume', type=int,
                        help='The program resume point')
    parser.add_argument('--SeriesID', action='store', dest='SeriesID',
                        help='The program series id')
    parser.add_argument('--StartTime', action='store', type=datetimeCheck, dest='StartTime',
                        help='The program start time')
    parser.add_argument('--Synopsis', action='store', dest='Synopsis',
                        help='The program synopsis')
    parser.add_argument('--Title', action='store', dest='Title',
                        help='The program title')
    parser.add_argument('--Team1', action='store', dest='Team1',
                        help='The first team (sports)')
    parser.add_argument('--Team2', action='store', dest='Team2',
                        help='The second team (sports)')

    args = parser.parse_args()

    # Build metadata from args
    metadata = {}
    for m in ['Category', 'ChannelAffiliate', 'ChannelImageURL', 'ChannelName', 'ChannelNumber',
              'EndTime', 'EpisodeNumber', 'EpisodeTitle', 'FirstAiring', 'ImageURL',
              'OriginalAirdate', 'ProgramID', 'RecordEndTime', 'RecordStartTime', 'RecordSuccess',
              'Resume', 'SeriesID', 'StartTime', 'Synopsis', 'Team1', 'Team2', 'Title']:
        value = vars(args)[m]
        if value is not None:
            metadata[m] = value

    # Open file
    openmode = 'rb'
    if args.set or args.update or args.delete:
        openmode = 'r+b'
    try:
        ffile = open(args.file, openmode)
    except (OSError, IOError) as e:
        print('Unable to open input file: ' + str(e))
        sys.exit(1)

    # Check metadata at front
    ffile.seek(0, 0)
    pkts = ffile.read(12032)
    try:
        fileMetadata = transportStreamToMetadata(pkts)
    except (ValueError, TypeError) as e:
        print('File {} does not appear to contain valid metadata at the start of the file: {}'.format(args.file, str(e)))
        sys.exit(1)

    # Perform the request
    if (not args.set) and (not args.update) and (not args.delete):
        if metadata:
            print('WARNING: Metadata values specified ignored, did you mean to --update or --set?')
        print(json.dumps(fileMetadata, sort_keys=True, indent=4))
    else:
        if args.update:
            fileMetadata.update(metadata)
        elif args.set:
            fileMetadata = metadata
        elif args.delete:
            for key in args.delete:
                try:
                    del fileMetadata[key]
                except KeyError:
                    pass
        try:
            ts = metadataToTransportStream(fileMetadata)
        except (ValueError, TypeError) as e:
            print('Unable to convert specified metadata to transport stream: {}'.format(str(e)))
            sys.exit(1)
        ffile.seek(0, 0)
        ffile.write(ts)
        ffile.seek(0, 2)

    ffile.close()




# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4