From bdc3b4cfc5c7e402a1c8148f2e4a27a9bbcdd264 Mon Sep 17 00:00:00 2001
From: Saxon <saxon@ausocean.org>
Date: Wed, 24 Jul 2019 00:04:10 +0930
Subject: [PATCH] container/mts/mpegts.go: added Programs, Streams and
 MediaStreams functions

---
 container/mts/mpegts.go | 85 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 82 insertions(+), 3 deletions(-)

diff --git a/container/mts/mpegts.go b/container/mts/mpegts.go
index 0b356770..f2a9ae0b 100644
--- a/container/mts/mpegts.go
+++ b/container/mts/mpegts.go
@@ -33,6 +33,7 @@ import (
 	"fmt"
 
 	"github.com/Comcast/gots/packet"
+	gotspsi "github.com/Comcast/gots/psi"
 	"github.com/pkg/errors"
 
 	"bitbucket.org/ausocean/av/container/mts/meta"
@@ -420,9 +421,9 @@ func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
 }
 
 var (
-	errNoPesPayload = errors.New("no PES payload")
-	errNoPesPTS = errors.New("no PES PTS")
-	errInvalidPesHeader = errors.New("invalid PES header")
+	errNoPesPayload      = errors.New("no PES payload")
+	errNoPesPTS          = errors.New("no PES PTS")
+	errInvalidPesHeader  = errors.New("invalid PES header")
 	errInvalidPesPayload = errors.New("invalid PES payload")
 )
 
@@ -592,3 +593,81 @@ func SegmentForMeta(d []byte, key, val string) ([][]byte, error) {
 
 	return res, nil
 }
+
+// pid returns the packet identifier for the given packet.
+func pid(p []byte) uint16 {
+	return uint16(p[1]&0x1f)<<8 | uint16(p[2])
+}
+
+// Programs returns a map of program numbers and corresponding PMT PIDs for a
+// given MPEG-TS PAT packet.
+func Programs(p []byte) (map[uint16]uint16, error) {
+	pat, err := gotspsi.NewPAT(p)
+	if err != nil {
+		return nil, err
+	}
+	return pat.ProgramMap(), nil
+}
+
+// Streams returns elementary streams defined in a given MPEG-TS PMT packet.
+// A gotspsi.PmtElementaryStream will give stream type from
+// gotspsi.PmtElementaryStream.StreamType() and PID from
+// gotspsi.PmtElementaryStream.ElementaryPid().
+//
+// PmtStreamTypes from gots/psi are defined as follows:
+// PmtStreamTypeMpeg2VideoH262 uint8 = 2  // H262
+// PmtStreamTypeMpeg4Video     uint8 = 27 // H264
+// PmtStreamTypeMpeg4VideoH264 uint8 = 27 // H264
+// PmtStreamTypeMpeg4VideoH265 uint8 = 36 // H265
+// PmtStreamTypeAac uint8 = 15  // AAC
+// PmtStreamTypeAc3 uint8 = 129 // DD
+// PmtStreamTypeEc3 uint8 = 135 // DD+
+// PmtStreamTypeScte35 uint8 = 134 // SCTE-35
+func Streams(p []byte) ([]gotspsi.PmtElementaryStream, error) {
+	pmt, err := gotspsi.NewPMT(p)
+	if err != nil {
+		return nil, err
+	}
+	return pmt.ElementaryStreams(), nil
+}
+
+// MediaStreams retrieves the PmtElementaryStreams from the given PSI. This
+// function currently assumes that PSI contain a PAT followed by a PMT directly
+// after. We also assume that this MPEG-TS stream contains just one program,
+// but this program may contain different streams, i.e. a video stream + audio
+// stream.
+func MediaStreams(p []byte) ([]gotspsi.PmtElementaryStream, error) {
+	pat := p[:PacketSize]
+	pmt := p[PacketSize : 2*PacketSize]
+
+	if pid(pat) != PatPid {
+		return nil, errors.New("first packet is not a PAT")
+	}
+
+	m, err := Programs(pat)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get programs from PAT")
+	}
+
+	if len(m) == 0 {
+		return nil, errors.New("no programs contained in PAT")
+	}
+
+	if len(m) > 1 {
+		return nil, errors.New("more than one program not yet supported")
+	}
+
+	var v uint16
+	for _, v = range m {
+	}
+
+	if pid(pmt) != v {
+		return nil, errors.New("second packet is not desired PMT")
+	}
+
+	s, err := Streams(pmt)
+	if err != nil {
+		return nil, errors.Wrap(err, "could not get streams from PMT")
+	}
+	return s, nil
+}