sppspatch

#!/bin/sh -efu

# v=0
# m=video <multicast_port> RTP/AVP 96
# c=IN IP4 <multicast_ip>
# a=rtpmap:96 H264/90000
# a=fmtp:96 sprop-parameter-sets=Z0IAKeKQFgJNgScFAQXh4kRU,aM48gA==
# a=source-filter: incl IN IP4 <multicast_ip> <unicast_ip>

PROG="${0##*/}"
PROG_VERSION=0.1.0

[ -z "${DEBUG:-}" ] || set -x

show_help() {
    cat <<EOF
Usage: $PROG [options] [-s <SDP-file>] [<inputfile>]

$PROG patches the given h264 stream with SPS/PPS info taken from a
SDP file. It also can be used to view the decoded SPS/PPS info.

Options:

   -g, --get              show decoded SPS/PPS; the SDP data is
                          read from <inputfile>, the -s argument or
                          stdin;

   -p, --patch            patch the given h264 stream file (the
                          default); if no inputfile is given then
                          stdin is used;

   -E, --no-escapes       do not insert the Emulation Prevention bytes
                          (03) in NALU sequences;

   -i, --inplace          patch the h264 stream file inplace;
                          otherwise the patched stream is written to
                          stdout;

   -s FILE, --sdp=FILE    the file to read SDP data from;

   -v, --verbose          be verbose;

   -V, --version          print version information and exit;

   -h, --help             print this help.
EOF
    exit 0
}

show_usage() {
    echo
    show_help | head -1
    echo "Use \`$PROG --help\` for more information."
}

print_version()
{
    cat <<EOF
$PROG version $PROG_VERSION
EOF
}

write_error() {
    printf "$@" >&2
}

info() {
    [ -z "$verbose" ] || write_error "$@"
}

fatal() {
    write_error "$@"
    exit 1
}


#####################
# Arguments parsing #
#####################

OPTS=`getopt -n $PROG -o v,V,h,s:,g,p,i,E -l verbose,version,help,sdp:,get,patch,inplace,no-escapes -- "$@"` || ( ret=$?; show_usage; exit $ret ) >&2
eval set -- "$OPTS"

verbose=; get_mode=; patch_mode=-p; inplace=; sdp=; escape=-e
while :; do
    case "$1" in
    -v|--verbose) verbose=-v;;
    -V|--version) print_version; exit 0;;
    -h|--help) show_help;;
    -s|--sdp) sdp="$2"; shift;;
    -g|--get)
            get_mode=-g
            patch_mode=
            ;;
    -E|--no-escapes) escape=;;
    -p|--patch)
            patch_mode=-p
            get_mode=
            ;;
    -i|--inplace) inplace=-i;;
    --) shift; break;;
    *)
        fatal 'Unrecognized option: %s\n' "$1"
        ;;
    esac
    shift
done


## Functions

getspps()
{
    local sdp="${1:-}"

    sed -n -e 's/^\(.\+[,;[:space:]]\+\)\?sprop-parameter-sets=\([^[:space:]]*\)[[:space:]]*;\?$/\2/p' \
        ${sdp:+"$sdp"}
}

getsps()
{
    sed -e 's/,.*$//'
}

getpps()
{
    sed -e 's/^.*,//'
}

# TODO: use bit mask
checkfirstbyte()
{
    local file="$1"
    local ref="${2:-}"

    if [ -z "$ref" ]; then
        ref="$file"
        file=
    fi
    
    local byte="$(hexdump -n 1 -e '/1 "%02X"' ${file:+"$file"})"

    [ "$byte" = "$ref" ]
}

writespps()
{
    local sdp="$1"

    info 'Read SDP data from %s\n' "${sdp:-<stdin>}"

    local spps="$(getspps "$sdp")"
    [ -n "$spps" ] || fatal 'SPS/PPS data not found!\n'

    local ret=0

    local sps="$(echo "$spps" | getsps)"
    if [ -n "$sps" ]; then
        echo "$sps" | base64 -d >"$workdir/sps" 2>/dev/null || ret=$?
        [ $ret -eq 0 ] || fatal 'Invalid SPS data: %s\n' "$sps"
        
        #if checkfirstbyte "$workdir/sps" 68; then
        #    mv "$workdir/sps" "$workdir/pps"
        #fi
    fi

    local pps="$(echo "$spps" | getpps)"
    if [ -n "$pps" ]; then
        echo "$pps" | base64 -d >"$workdir/pps" 2>/dev/null || ret=$?
        [ $ret -eq 0 ] || fatal 'Invalid PPS data: %s\n' "$pps"

        #if checkfirstbyte "$workdir/pps" 67; then
        #    mv "$workdir/pps" "$workdir/sps"
        #fi
    fi
}

printinhex()
{
    local prefix="$1"
    local file="${2:-}"
    local zcnt=0

    if [ -z "$file" ]; then
        file="$prefix"
        prefix=
    fi

    [ -s "$file" ] || return 0
    
    if [ -z "$prefix" ]; then
        cat "$file" | hexdump -v -e '/1 "%02X "'
    elif [ -z "$escape" ]; then
        printf '%s' "$prefix"
        cat "$file" | hexdump -v -e '/1 " %02X"'
    else
        printf '%s' "$prefix"
        cat "$file" | hexdump -v -e '/1 "%02X\n"' | \
            while read byte; do
                case "$byte" in
                    00|01|02)
                        if [ $zcnt -eq 2 ]; then
                            printf ' 03'
                            zcnt=0
                        elif [ "$byte" = '00' ]; then
                            zcnt=$((zcnt + 1))
                        fi
                        printf ' %s' "$byte"
                        ;;
                    *)
                        printf ' %s' "$byte"
                        zcnt=0
                    ;;
                esac
            done
    fi
}

hexbytes()
{
    (
        tr ' ' '\n'
        echo
    ) | while read hex; do
            [ -z "$hex" ] || /usr/bin/printf "\\x$hex"
        done
}

makenalu()
{
    local file="$1"
    local prefix='00 00 00 01'

    printinhex "$prefix" "$file"
}


## Main

workdir="$(mktemp --tmpdir -d "$PROG.XXXXXXXXXX")"

cleanup()
{
    [ -z "$workdir" ] || rm -rf "$workdir"
}
trap "cleanup" EXIT


if [ -n "$get_mode" ]; then
    # Get mode
    [ -n "$sdp" ] || sdp="${1:-}"
    writespps "$sdp"
    [ ! -e "$workdir/sps" ] || \
        printf 'SPS: %s\n' "$(printinhex "$workdir/sps")"
    [ ! -e "$workdir/pps" ] || \
        printf 'PPS: %s\n' "$(printinhex "$workdir/pps")"
else
    # Patch mode
    infile="${1:-}"
    [ -n "$sdp" ] || fatal 'SDP file is not specified\n'
    
    writespps "$sdp"

    [ -e "$workdir/sps" ] || fatal 'SPS data not found\n'
    info 'SPS: %s\n' "$(printinhex "$workdir/sps")"

    [ -e "$workdir/pps" ] || fatal 'PPS data not found\n'
    info 'PPS: %s\n' "$(printinhex "$workdir/pps")"

    info 'Read h264 stream from %s\n' "${infile:-<stdin>}"

    info 'Escaping is %s\n' \
        "$(if [ -n "$escape" ]; then echo 'on'; else echo 'off'; fi)"

    makenalu "$workdir/sps" >"$workdir/spsnalu"
    info 'SPS NALU: %s\n' "$(cat "$workdir/spsnalu")"

    makenalu "$workdir/pps" >"$workdir/ppsnalu"
    info 'PPS NALU: %s\n' "$(cat "$workdir/ppsnalu")"
    
    if [ -z "$inplace" -o -z "$infile" ]; then
        cat "$workdir/spsnalu" | hexbytes
        cat "$workdir/ppsnalu" | hexbytes
        cat "${infile:--}"
    else
        info 'Patching the stream inplace\n'
        (
            cat "$workdir/spsnalu" | hexbytes
            cat "$workdir/ppsnalu" | hexbytes
            cat "$infile"
        ) >"$workdir/outfile"
        mv "$workdir/outfile" "$infile"
    fi
fi