cmd/ts-repair: got setting of discontinuity indicators working and also adding adaptation fields to pat and pmt for this reason.

cmd/ts-repair/main.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"flag"
+	"fmt"
 	"io"
 	"os"
 
@@ -15,25 +16,94 @@ const (
 	errCantGetPid    = "Can't get pid from packet"
 	errReadFail      = "Read failed"
 	errWriteFail     = "Write to file failed"
-	usage            = "The path to the file to be repaired"
+	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:   1,
+	mts.PatPid:   16,
-	mts.PmtPid:   1,
+	mts.PmtPid:   16,
-	mts.VideoPid: 0,
+	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("path", "", usage)
+	inPtr := flag.String("in", "", inUsage)
-	outPtr := flag.String("out", "out.ts", usage)
+	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
@@ -59,14 +129,42 @@ func main() {
 		} 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)
 		}
-		p.setCC(ccFor(int(pid)))
 
+		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("packetNo: %v pid: %v, cc: %v, expect: %v\n", packetNo, pid, cc, expect)
+					if p.hasAdaptation() {
+						fmt.Println("hasAdaptation")
+						p.setDI(true)
+					} else {
+						fmt.Println("doesn't have adaptation")
+						fmt.Println(p)
+						p.addAdaptationField(DiscontinuityIndicator(true))
+						fmt.Println(p)
+					}
+					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())
@@ -75,8 +173,15 @@ func main() {
 }
 
 // ccFor gets the next cc for the given pid
-func ccFor(pid int) byte {
+func expectedCC(pid int) (byte, bool) {
 	cc := ccMap[pid]
-	ccMap[pid] = (cc + 1) & 0xf
+	if cc == 16 {
-	return cc
+		return 16, false
+	}
+	ccMap[pid] = (cc + 1) & 0xf
+	return cc, true
+}
+
+func updateCCMap(pid int, cc byte) {
+	ccMap[pid] = (cc + 1) & 0xf
 }