diff --git a/cmd/ts-repair/main.go b/cmd/ts-repair/main.go index bcb1ebc1..f34da350 100644 --- a/cmd/ts-repair/main.go +++ b/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.PmtPid: 1, - mts.VideoPid: 0, + 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("path", "", usage) - outPtr := flag.String("out", "out.ts", usage) + 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 @@ -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] + 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 - return cc }