Select a program stream in a Transport Stream.

tsslice.c

/*
 * ***** BEGIN LICENSE BLOCK *****
 * Version: MPL 1.1
 *
 * The contents of this file are subject to the Mozilla Public License Version
 * 1.1 (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 * http://www.mozilla.org/MPL/
 *
 * Software distributed under the License is distributed on an "AS IS" basis,
 * WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
 * for the specific language governing rights and limitations under the
 * License.
 *
 * The Original Code is the MPEG TS, PS and ES tools.
 *
 * The Initial Developer of the Original Code is Amino Communications Ltd.
 * Portions created by the Initial Developer are Copyright (C) 2008
 * the Initial Developer. All Rights Reserved.
 *
 * Contributor(s):
 *   Amino Communications Ltd, Swavesey, Cambridge UK
 *
 * ***** END LICENSE BLOCK *****
 */

#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>

#ifdef _WIN32
#include <stddef.h>
#else // _WIN32
#include <unistd.h>
#endif // _WIN32

#include "compat.h"
#include "ts_fns.h"
#include "misc_fns.h"
#include "tswrite_fns.h"
#include "pidint_fns.h"
#include "version.h"


static void filter_streams(pmt_p pmt)
{
    static const uint32_t av_streams[] = 
	{
	    0x01, // return "11172-2 video (MPEG-1)";
	    0x02, // return "H.262/13818-2 video (MPEG-2) or 11172-2 constrained video";
	    0x03, // return "11172-3 audio (MPEG-1)";
		    0x04, // return "13818-3 audio (MPEG-2)";
	    0x07, // return "13522 MHEG";
	    0x08, // return "H.222.0/13818-1 Annex A - DSM CC";
	    0x09, // return "H.222.1";
	    0x0F, // return "13818-7 Audio with ADTS transport syntax";
	    0x10, // return "14496-2 Visual (MPEG-4 part 2 video)";
	    0x11, // return "14496-3 Audio with LATM transport syntax (14496-3/AMD 1)";
	    0x1B, // return "H.264/14496-10 video (MPEG-4/AVC)";
	};
    int i, pids[pmt->num_streams], n;
    for (n = i = 0; i < pmt->num_streams; i++) {
	int j;
	for (j = 0; j < sizeof(av_streams) / sizeof(av_streams[0]); j++)
	  if (pmt->streams[i].stream_type == av_streams[j])
	    break;
	if (j >= sizeof(av_streams) / sizeof(av_streams[0]))
	  pids[n++] = pmt->streams[i].elementary_PID;
    }
    for (i = 0; i < n; i++)
      remove_stream_from_pmt(pmt, pids[i]);
}


/*
 * Report on the program streams, by looking at the PAT and PMT packets
 * in the first `max` TS packets of the given input stream
 *
 * Returns 0 if all went well, 1 if something went wrong.
 */
static int select_stream(TS_reader_p tsreader, 
			 TS_writer_p  tswriter,
			 int program, int ordinal)
{
    char    pid_ar[8192];
    int     err;
    int     ii;

    // TODO: Should really support multiple programs
    //       (some use of pidint_list to support program number -> PMT?)

    pidint_list_p  this_prog_list = NULL;
    pidint_list_p  last_prog_list = NULL;
    pmt_p          this_pmt = NULL;
    pmt_p          last_pmt = NULL;

    uint32_t       pmt_pid = 0;  // which will get "masked" by the PAT pid
    int            program_id = 0;

    byte  *pat_data = NULL;
    int    pat_data_len = 0;
    int    pat_data_used = 0;

    byte  *pmt_data = NULL;
    int    pmt_data_len = 0;
    int    pmt_data_used = 0;

    int num_pats = 0;
    int num_pmts = 0;

    memset(pid_ar, 0, sizeof(pid_ar));
    for (ii = 0; 1 == 1; ii++) {
	uint32_t pid;
	int      payload_unit_start_indicator;
	byte    *adapt, *payload;
	int      adapt_len, payload_len;
	byte  *packet;

	err = read_next_TS_packet(tsreader, &packet);
	if (err == EOF)
	  break;
	else if (err) {
	    fprintf(stderr,"### Error reading TS packet\n");
	    return 1;
	}
	err = split_TS_packet(packet, &pid, &payload_unit_start_indicator,
			      &adapt, &adapt_len, &payload, &payload_len);
	if (err == EOF)
	  break;
	else if (err) {
	    fprintf(stderr,"### Error reading TS packet %d\n",ii+1);
	    if (pat_data) free(pat_data);
	    free_pidint_list(&last_prog_list);
	    free_pmt(&last_pmt);
	    if (pmt_data) free(pmt_data);
	    return 1;
	}

	if (pid == 0x0000) {
	    num_pats++;
	    if (payload_len == 0)  {
		fprintf(stderr, "Packet %d is PAT, but has no payload\n",ii+1);
		continue;
	    }

	    if (payload_unit_start_indicator && pat_data)  {
		// This is the start of a new PAT packet, but we'd already
		// started one, so throw its data away
		fprintf(stderr,"!!! Discarding previous (uncompleted) PAT data\n");
		free(pat_data);
		pat_data = NULL; pat_data_len = 0; pat_data_used = 0;
	    } else if (!payload_unit_start_indicator && !pat_data) {
		// This is the continuation of a PAT packet, but we hadn't
		// started one yet
		fprintf(stderr,"!!! Discarding PAT continuation, no PAT started\n");
		continue;
	    }

	    err = build_psi_data(0, payload,payload_len,pid,
				 &pat_data,&pat_data_len,&pat_data_used);
	    if (err) {
		fprintf(stderr,"### Error %s PAT\n",
			(payload_unit_start_indicator?"starting new":"continuing"));
		if (pat_data) free(pat_data);
		free_pidint_list(&last_prog_list);
		free_pmt(&last_pmt);
		if (pmt_data) free(pmt_data);
		return 1;
	    }

	    // Do we need more data to complete this PAT?
	    if (pat_data_len > pat_data_used)
	      continue;

	    err = extract_prog_list_from_pat(0, pat_data,pat_data_len,
					     &this_prog_list);
	    if (err) {
		free_pidint_list(&last_prog_list);
		free_pmt(&last_pmt);
		if (pmt_data) free(pmt_data);
		free(pat_data);
		return err;
	    }

	    free(pat_data);
	    pat_data = NULL; pat_data_len = 0; pat_data_used = 0;
	    num_pats++;

	    if (err) {
		free_pidint_list(&last_prog_list);
		free_pmt(&last_pmt);
		if (pmt_data) free(pmt_data);
		return err;
	    }

	    if (!same_pidint_list(this_prog_list, last_prog_list)) {
		memset(pid_ar, 0, sizeof(pid_ar));
		pmt_pid = 0;

		if (this_prog_list->length > 0) {
		    int i;
		    if (program < 0) {
			for (i = 0; i < this_prog_list->length; i++) {
			    if (this_prog_list->pid[i] == -program) {
				pmt_pid = this_prog_list->pid[i];
				program_id = this_prog_list->number[i];
				break;
			    }
			}
			if (pmt_pid == 0) {
			    fprintf(stderr, "PMT_PID = %x does not found. skip PAT.\n", -program);
			    continue;
			}
		    } else if (ordinal) {
			if (program < this_prog_list->length) {
			    pmt_pid = this_prog_list->pid[program];
			    program_id = this_prog_list->number[program];
			} else {
			    fprintf(stderr, "program #%d does not found. skip PAT.\n", program);
			    continue;
			}
		    } else {
			for (i = 0; i < this_prog_list->length; i++) {
			    if (this_prog_list->number[i] == program) {
				pmt_pid = this_prog_list->pid[i];
				program_id = this_prog_list->number[i];
				break;
			    }
			}
			if (pmt_pid == 0) {
			    fprintf(stderr, "PMT_PID = %x does not found. skip PAT.\n", program);
			    continue;
			}
		    }
		}
		fprintf(stderr, "new PMT %x\n", pmt_pid);
	    }

	    {
		struct pidint_list pat = {
		    &program_id,
		    &pmt_pid,
		    1, 1,
		};
		int transport_stream_id = (packet[3] << 8) | packet[4];
		err = write_pat(tswriter, transport_stream_id,  &pat);
		if (err)
		  return 1;
	    }

	    free_pidint_list(&last_prog_list);
	    last_prog_list = this_prog_list;
	} else if (pid == pmt_pid) {
	    if (payload_len == 0) {
		fprintf(stderr, "Packet %d is PMT, but has no payload\n",ii+1);
		continue;
	    }

	    if (payload_unit_start_indicator && pmt_data) {
		// This is the start of a new PMT packet, but we'd already
		// started one, so throw its data away
		fprintf(stderr,"!!! Discarding previous (uncompleted) PMT data\n");
		free(pmt_data);
		pmt_data = NULL; pmt_data_len = 0; pmt_data_used = 0;
	    } else if (!payload_unit_start_indicator && !pmt_data) {
		// This is the continuation of a PMT packet, but we hadn't
		// started one yet
		fprintf(stderr,"!!! Discarding PMT continuation, no PMT started\n");
		continue;
	    }

	    err = build_psi_data(0, payload,payload_len,pid,
				 &pmt_data,&pmt_data_len,&pmt_data_used);
	    if (err) {
		fprintf(stderr,"### Error %s PMT\n",
			(payload_unit_start_indicator?"starting new":"continuing"));
		free_pidint_list(&this_prog_list);
		free_pmt(&last_pmt);
		if (pmt_data) free(pmt_data);
		return 1;
	    }

	    // Do we need more data to complete this PMT?
	    if (pmt_data_len > pmt_data_used)
	      continue;

	    err = extract_pmt(0, pmt_data, pmt_data_len,pid, &this_pmt);
	    if (err) {
		free_pidint_list(&this_prog_list);
		free_pmt(&last_pmt);
		if (pmt_data) free(pmt_data);
		return err;
	    }

	    free(pmt_data);
	    pmt_data = NULL; pmt_data_len = 0; pmt_data_used = 0;
	    num_pmts++;

	    filter_streams(this_pmt);

	    err = write_pmt(tswriter, pmt_pid, this_pmt);
	    if (err)
	      return 1;

	    if (same_pmt(this_pmt,last_pmt)) { // Nothing to do
		free_pmt(&this_pmt);
		continue;
	    }
	    
	    fprintf(stderr, "PIDS: %x(PCR)", this_pmt->PCR_pid);
	    pid_ar[this_pmt->PCR_pid] = 1;
	    pid_ar[0x14] = 1;
	    
	    {
		int i;
		for (i=0; i < this_pmt->num_streams; i++) {
		    pid_ar[this_pmt->streams[i].elementary_PID] = 1;
		    fprintf(stderr, " %x(%x)", 
			    this_pmt->streams[i].elementary_PID, 
			    this_pmt->streams[i].stream_type);
		}
		putc('\n', stderr);
	    }
	    free_pmt(&last_pmt);
	    last_pmt = this_pmt;
	} else if (pid_ar[pid]) {
	    err = tswrite_write(tswriter, packet, 0, FALSE, 0);
	    if (err) 
	      return 1;
	}
    }

    fprintf(stderr, "\nFound %d PAT packet%s and %d PMT packet%s in %d TS packets\n",
	    num_pats,(num_pats==1?"":"s"),
	    num_pmts,(num_pmts==1?"":"s"), ii);

    free_pidint_list(&last_prog_list);
    free_pmt(&last_pmt);
    if (pmt_data) free(pmt_data);
    return 0;
}


static void print_usage()
{
    printf("Usage: tsselect [switches] [<infile>]\n"
	   "Select a program stream in a Transport Stream.\n"
	   "\n"
	   "Files:\n"
	   "  <infile>  is an H.222 Transport Stream file (but see -stdin)\n"
	   "\n"
	   "Switches:\n"
	   " -stdin			Input from standard input, instead of a file\n"
	   " -program <n>, -p <n>	index of Program\n"
	   " -pmt-pid <n>, -P <n>	PID of PMT\n");
}

int main(int argc, char **argv)
{
    char  *input_name = NULL;
    char  *output_name = NULL;
    int    err = 0;
    int	   program = 0;
    int    ordinal = 0;

    TS_reader_p  tsreader = NULL;
    TS_writer_p  tswriter = NULL;

    int    ii = 1;

    if (argc < 2)
    {
	print_usage();
	return 0;
    }

    while (ii < argc) {
	if (argv[ii][0] == '-' && argv[ii][1] != 0) {
	    if (!strcmp("--help",argv[ii]) || !strcmp("-h",argv[ii]) ||
		!strcmp("-help",argv[ii]))
	    {
		print_usage();
		return 0;
	    } else if (!strcmp("-program",argv[ii]) || !strcmp("-p",argv[ii])) {
		CHECKARG("tsinfo",ii);
		err = int_value("tsinfo",argv[ii],argv[ii+1],TRUE, 0,&program);
		if (err) 
		  return 1;
		if (argv[ii+1][0] == '+')
		  ordinal = 1;
		ii++;
	    } else if (!strcmp("-pmt_pid",argv[ii]) || !strcmp("-P",argv[ii])) {
		CHECKARG("tsinfo",ii);
		err = int_value("tsinfo",argv[ii],argv[ii+1],TRUE, 0,&program);
		program = -program;
		if (err) return 1;
		ii++;
	    } else {
		fprintf(stderr,"### tsselect: "
			"Unrecognised command line switch '%s'\n",argv[ii]);
		return 1;
	    }
	} else {
	    if (input_name == NULL)
	      input_name = argv[ii];
	    else if (output_name == NULL)
	      output_name = argv[ii];
	    else {
		fprintf(stderr,"### tsselect: Unexpected '%s'\n",argv[ii]);
		return 1;
	    }
	}
	ii++;
    }

    if (input_name == NULL) {
	fprintf(stderr,"### tsselect: No input file specified\n");
	return 1;
    }
    if (output_name == NULL) {
	fprintf(stderr,"### tsselect: No output file specified\n");
	return 1;
    }

    err = open_file_for_TS_read(strcmp("-", input_name) ? input_name : NULL ,&tsreader);
    if (err) {
	fprintf(stderr,"### tsselect: Unable to open input file %s for reading TS\n",
		strcmp("-", input_name) ? input_name : "<stdin>");
	return 1;
    }

    err = tswrite_open(strcmp("-", output_name)? TS_W_FILE : TS_W_STDOUT,
		       output_name, NULL,  -1, TRUE, &tswriter);

    if (err) {
	fprintf(stderr,"### tsselect: Unable to open output file %s for writing TS\n",
		strcmp("-", output_name) ? output_name : "<stdout>");
	return 1;
    }

    err = select_stream(tsreader, tswriter, program, ordinal);

    if (err) {
	fprintf(stderr,"### tsselect: Error reporting on stream\n");
	(void) close_TS_reader(&tsreader);
	return 1;
    }
    tswrite_close(tswriter, TRUE);

    err = close_TS_reader(&tsreader);
    if (err) return 1;

    return 0;
}