arib-cli 0.1.0

Reads the signalling of ARIB broadcasts and descrambles them, as an example of the arib crate
use std::fs::File;
use std::io::{BufRead, BufReader, Read, stdin};
use std::path::PathBuf;

use anyhow::{Context, bail};
use clap::{Args, ValueEnum};

/// Where a stream is read from, and what it is.
#[derive(Clone, Debug, Args)]
pub struct InputArgs {
    /// Path of the stream to read, or `-` for stdin.
    #[arg(default_value = "-")]
    input: PathBuf,

    /// What the stream is, told by its first bytes unless given.
    #[arg(long, value_enum, default_value_t = FormatArg::Auto)]
    format: FormatArg,
}

#[derive(Copy, Clone, Debug, ValueEnum)]
enum FormatArg {
    Auto,
    /// MPEG-2 TS, as ISDB-T and ISDB-S broadcast.
    Ts,
    /// MMT over TLV, as ISDB-S3 broadcasts.
    Mmt,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum Format {
    Ts,
    Mmt,
}

pub struct Input {
    pub reader: BufReader<Box<dyn Read>>,
    pub format: Format,
}

impl InputArgs {
    pub fn open(&self) -> anyhow::Result<Input> {
        let reader: Box<dyn Read> = if self.input.as_os_str() == "-" {
            Box::new(stdin())
        } else {
            Box::new(File::open(&self.input).with_context(|| format!("{}", self.input.display()))?)
        };
        let mut reader = BufReader::with_capacity(1 << 20, reader);

        let format = match self.format {
            FormatArg::Ts => Format::Ts,
            FormatArg::Mmt => Format::Mmt,
            FormatArg::Auto => match detect(reader.fill_buf()?) {
                Some(format) => format,
                None => bail!("could not tell what the stream is; give it with --format"),
            },
        };

        Ok(Input { reader, format })
    }
}

/// Tells the format by three packets in a row: of 188 bytes starting with the sync byte, or of
/// TLV, each saying how long it is.
fn detect(buf: &[u8]) -> Option<Format> {
    const TS_SIZE: usize = 188;

    let is_ts =
        |offset: usize| (0..3).all(|index| buf.get(offset + index * TS_SIZE) == Some(&0x47));
    if (0..TS_SIZE).any(is_ts) {
        return Some(Format::Ts);
    }

    let next_tlv = |offset: usize| {
        let head = buf.get(offset..offset + 4)?;
        let known_type = matches!(head[1], 0x01 | 0x02 | 0x03 | 0xFE | 0xFF);
        (head[0] == 0x7F && known_type)
            .then(|| offset + 4 + usize::from(u16::from_be_bytes([head[2], head[3]])))
    };
    let is_tlv = |offset: usize| {
        next_tlv(offset)
            .and_then(next_tlv)
            .and_then(next_tlv)
            .is_some()
    };
    (0..buf.len().min(1 << 16))
        .any(is_tlv)
        .then_some(Format::Mmt)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn tells_ts_by_its_sync_bytes() {
        let mut stream = vec![0x00; 10];
        for _ in 0..3 {
            stream.extend([0x47].iter().chain(&[0xFF; 187]));
        }

        assert_eq!(detect(&stream), Some(Format::Ts));
    }

    #[test]
    fn tells_tlv_by_the_lengths_of_its_packets() {
        let packet = [0x7F, 0xFF, 0x00, 0x02, 0xFF, 0xFF];
        let stream = [&[0x12, 0x34][..], &packet, &packet, &packet].concat();

        assert_eq!(detect(&stream), Some(Format::Mmt));
        assert_eq!(detect(&[0x12; 1024]), None);
    }
}