av/cmd/ts-repair/main.go

83 lines
1.8 KiB
Go

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
}