av/cmd/ts-repair/main.go

package main

import (
	"flag"
	"fmt"
	"io"
	"os"

	"bitbucket.org/ausocean/av/stream/mts"
	"github.com/Comcast/gots/packet"
)

const (
	errBadInPath     = "No file path provided, or file does not exist"
	errCantCreateOut = "Can't create output file"
	errCantGetPid    = "Can't get pid from packet"
	errReadFail      = "Read failed"
	errWriteFail     = "Write to file failed"
	errBadMode       = "Bad fix mode"
)

const (
	inUsage   = "The path to the file to be repaired"
	outUsage  = "Output file path"
	modeUsage = "Fix mode: 0 = cc-shift, 1 = di-update"
)

const (
	ccShift = iota
	diUpdate
)

var ccMap = map[int]byte{
	mts.PatPid:   16,
	mts.PmtPid:   16,
	mts.VideoPid: 16,
}

var packetNo int

type Option func(p *Packet)

type Packet [mts.PacketSize]byte

func (p *Packet) CC() byte {
	return (*p)[3] & 0x0f
}

func (p *Packet) setCC(cc byte) {
	(*p)[3] |= cc & 0xf
}

func (p *Packet) setDI(di bool) {
	if di {
		p[5] |= 0x80
	} else {
		p[5] &= 0x7f
	}
}

func (p *Packet) addAdaptationField(options ...Option) {
	// Create space for adaptation field
	copy(p[mts.HeadSize+mts.DefaultAdaptationSize:], p[mts.HeadSize:len(p)-mts.DefaultAdaptationSize])

	// TODO: seperate into own function
	// Update adaptation field control
	p[mts.AdaptationControlIdx] &= 0xff ^ mts.AdaptationControlMask
	p[mts.AdaptationControlIdx] |= mts.AdaptationControlMask
	// Default the adaptationfield
	p.resetAdaptation()

	for _, option := range options {
		option(p)
	}
}

func (p *Packet) resetAdaptation() {
	p[mts.AdaptationIdx] = mts.DefaultAdaptationBodySize
	p[mts.AdaptationBodyIdx] = 0x00
}

func (p *Packet) hasAdaptation() bool {
	afc := p[mts.AdaptationControlIdx] & mts.AdaptationControlMask
	if afc == 0x20 || afc == 0x30 {
		return true
	} else {
		return false
	}
}

func DiscontinuityIndicator(f bool) Option {
	return func(p *Packet) {
		set := byte(mts.DiscontinuityIndicatorMask)
		if !f {
			set = 0x00
		}
		p[mts.DiscontinuityIndicatorIdx] &= 0xff ^ mts.DiscontinuityIndicatorMask
		p[mts.DiscontinuityIndicatorIdx] |= mts.DiscontinuityIndicatorMask & set
	}
}

func main() {
	// Deal with input flags
	inPtr := flag.String("in", "", inUsage)
	outPtr := flag.String("out", "out.ts", outUsage)
	modePtr := flag.Int("mode", diUpdate, modeUsage)
	flag.Parse()

	// Try and open the given input file, otherwise panic - we can't do anything
	inFile, err := os.Open(*inPtr)
	defer inFile.Close()
	if err != nil {
		panic(errBadInPath)
	}

	// Try and create output file, otherwise panic - we can't do anything
	outFile, err := os.Create(*outPtr)
	defer outFile.Close()
	if err != nil {
		panic(errCantCreateOut)
	}

	// Read each packet from the input file reader
	var p Packet
	for {
		// If we get an end of file then return, otherwise we panic - can't do anything else
		if _, err := inFile.Read(p[:mts.PacketSize]); err == io.EOF {
			return
		} else if err != nil {
			panic(errReadFail + ": " + err.Error())
		}
		packetNo++
		// Get the pid from the packet and set the cc based on this pid using our map
		pid, err := packet.Pid((*packet.Packet)(&p))
		if err != nil {
			panic(errCantGetPid)
		}

		cc := p.CC()
		expect, exists := expectedCC(int(pid))
		if !exists {
			updateCCMap(int(pid), cc)
		} else {
			switch *modePtr {
			case ccShift:
				p.setCC(expect)

			case diUpdate:

				if cc != expect {
					fmt.Printf("***** Discontinuity found (packetNo: %v pid: %v, cc: %v, expect: %v)\n", packetNo, pid, cc, expect)
					if p.hasAdaptation() {
						p.setDI(true)
					} else {
						p.addAdaptationField(DiscontinuityIndicator(true))
					}
					updateCCMap(int(pid), p.CC())
				}

			default:
				panic(errBadMode)
			}
		}
		// Write this packet to the output file
		if _, err := outFile.Write(p[:]); err != nil {
			panic(errWriteFail + ": " + err.Error())
		}
	}
}

// ccFor gets the next cc for the given pid
func expectedCC(pid int) (byte, bool) {
	cc := ccMap[pid]
	if cc == 16 {
		return 16, false
	}
	ccMap[pid] = (cc + 1) & 0xf
	return cc, true
}

func updateCCMap(pid int, cc byte) {
	ccMap[pid] = (cc + 1) & 0xf
}