mts: fixing cc

cmd/ts-repair/main.go

@@ -0,0 +1,82 @@
+package main
+
+import (
+	"flag"
+	"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"
+	usage            = "The path to the file to be repaired"
+)
+
+var ccMap = map[int]byte{
+	mts.PatPid:   1,
+	mts.PmtPid:   1,
+	mts.VideoPid: 0,
+}
+
+type Packet [mts.PacketSize]byte
+
+func (p *Packet) setCC(cc byte) {
+	(*p)[3] |= cc & 0xf
+}
+
+func main() {
+	// Deal with input flags
+	inPtr := flag.String("path", "", usage)
+	outPtr := flag.String("out", "out.ts", usage)
+	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())
+		}
+
+		// 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)
+		}
+		p.setCC(ccFor(int(pid)))
+
+		// 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 ccFor(pid int) byte {
+	cc := ccMap[pid]
+	ccMap[pid] = (cc + 1) & 0xf
+	return cc
+}