arib-cli 0.2.0

use std::fs::File;
use std::io::{BufWriter, Write, stdout};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};

use anyhow::{Context, bail};
use arib::cas::pcsc::PcscCasModule;
use clap::Args;
use tracing::warn;

use crate::input::{Format, InputArgs};

/// Descrambles the stream with the CAS card in the first PC/SC reader, and writes it out.
#[derive(Clone, Debug, Args)]
pub struct Options {
    #[command(flatten)]
    input: InputArgs,

    /// Path to write the descrambled stream to, or `-` for stdout.
    #[arg(short, long, default_value = "-")]
    output: PathBuf,

    /// The master key Kd the ACAS card protects the keys with, in 64 hex digits. MMT only.
    #[arg(long, env = "ARIB_MASTER_KEY", hide_env_values = true, value_parser = parse_master_key)]
    master_key: Option<[u8; 32]>,

    /// Descrambles the service by the ECMs of its PMT only. MPEG-2 TS only.
    #[arg(long, value_parser = crate::print::parse_id)]
    service: Option<u16>,

    /// Goes on writing the stream while the card works on an ECM, as a stream received live, from
    /// a tuner, needs to. A recording is read far faster than it was broadcast, so the ECMs come
    /// faster than the card answers them and it is waited for instead unless this is given.
    #[arg(long)]
    live: bool,
}

pub fn run(options: &Options) -> anyhow::Result<()> {
    let input = options.input.open()?;
    // The command line is looked at before the card is, so that a mistake on it is what the
    // error talks about.
    if input.format == Format::Mmt && options.master_key.is_none() {
        bail!("MMT needs the master key: give it with --master-key or ARIB_MASTER_KEY");
    }

    let mut output: BufWriter<Box<dyn Write>> =
        BufWriter::new(if options.output.as_os_str() == "-" {
            Box::new(stdout())
        } else {
            Box::new(
                File::create(&options.output)
                    .with_context(|| format!("{}", options.output.display()))?,
            )
        });
    let module = Arc::new(Mutex::new(
        PcscCasModule::open().context("could not open the CAS card")?,
    ));

    match input.format {
        Format::Ts => {
            let descrambler = arib::ts::descramble::Descrambler::init(module, options.live)?;
            let mut demuxer =
                arib::ts::demux::Demuxer::new(input.reader).with_descrambler(descrambler);
            if let Some(service_id) = options.service {
                demuxer = demuxer.for_service(service_id);
            }

            while let Some(packet) = demuxer.read_packet()? {
                output.write_all(packet.as_bytes())?;
                demuxer.take_packets().try_for_each(report)?;
            }
        }
        Format::Mmt => {
            let master_key = options.master_key.expect("checked above");
            let descrambler =
                arib::mmt::descramble::Descrambler::init(module, master_key, options.live)?;
            let mut demuxer =
                arib::mmt::demux::Demuxer::new(input.reader).with_descrambler(descrambler);

            while let Some(packet) = demuxer.read_packet()? {
                packet.write(&mut output)?;
                demuxer.take_packets().try_for_each(report)?;
            }
        }
    }

    output.flush()?;

    Ok(())
}

/// Stops at a card refusing to descramble, which it will do to every ECM after; the rest is
/// only warned about, as the stream is still written out.
fn report(packet: arib::Result<arib::demux::Packet>) -> anyhow::Result<()> {
    match packet {
        Ok(_) => Ok(()),
        Err(error @ arib::Error::EcmRefused(_)) => Err(error.into()),
        Err(error) => {
            warn!(%error, "Could not descramble or read a packet");
            Ok(())
        }
    }
}

fn parse_master_key(value: &str) -> Result<[u8; 32], String> {
    let digits = value.trim();
    if digits.len() != 64 || !digits.is_ascii() {
        return Err("expected 64 hex digits".to_owned());
    }

    let mut key = [0u8; 32];
    for (byte, pair) in key.iter_mut().zip(digits.as_bytes().chunks(2)) {
        let pair = std::str::from_utf8(pair).expect("ASCII");
        *byte = u8::from_str_radix(pair, 16).map_err(|error| error.to_string())?;
    }

    Ok(key)
}

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

    #[test]
    fn parses_a_master_key_of_64_hex_digits() {
        let key = parse_master_key(&"0123456789abcdef".repeat(4)).unwrap();
        assert_eq!(key[..8], [0x01, 0x23, 0x45, 0x67, 0x89, 0xAB, 0xCD, 0xEF]);

        assert!(parse_master_key("0123").is_err());
        assert!(parse_master_key(&"xy".repeat(32)).is_err());
    }
}