view src/main.rs @ 1:4adea0a8c8dd

Remove all unwraps.
author Emmanuel Gil Peyrot <linkmauve@linkmauve.fr>
date Thu, 19 Nov 2020 21:36:00 +0100
parents 57b200fad67e
children 1a2852861d8b
line wrap: on
line source

use std::num::ParseIntError;
use std::str::FromStr;
use std::fmt;
use std::fs::File;
use std::io::{Read, BufReader, Seek, SeekFrom, Write, BufWriter};
use std::error::Error;
use std::env;
use std::process;

#[derive(Debug)]
pub enum ParseError {
    Int(ParseIntError),
    WrongResolution,
    UnknownCodec,
    Custom(String),
}

impl fmt::Display for ParseError {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}", self)
    }
}

impl Error for ParseError {}

impl From<ParseIntError> for ParseError {
    fn from(e: ParseIntError) -> ParseError {
        ParseError::Int(e)
    }
}

#[derive(Debug)]
enum Codec {
    H264,
}

impl FromStr for Codec {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Codec, ParseError> {
        if s == "H.264" {
            Ok(Codec::H264)
        } else {
            Err(ParseError::UnknownCodec)
        }
    }
}

#[derive(Debug, PartialEq)]
struct Resolution {
    width: u16,
    height: u16,
}

impl FromStr for Resolution {
    type Err = ParseError;

    fn from_str(s: &str) -> Result<Resolution, ParseError> {
        let values: Vec<_> = s.split('x').collect();
        if let [width, height] = values.as_slice() {
            let width = width.parse()?;
            let height = height.parse()?;
            Ok(Resolution {
                width,
                height,
            })
        } else {
            Err(ParseError::WrongResolution)
        }
    }
}

impl fmt::Display for Resolution {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{}×{}", self.width, self.height)
    }
}

#[derive(Debug)]
pub struct VxSeekLine {
    timestamp: u64,
    video_name: String,
    codec: Codec,
    resolution: Resolution,
    unknown: u8,
    length: u32,
}

pub fn parse_vxseek(vxseek: &str) -> Result<Vec<VxSeekLine>, ParseError> {
    let mut content = Vec::new();
    for line in vxseek.lines() {
        let fields: Vec<_> = line.split(':').collect();
        match fields.as_slice() {
            [timestamp, video_name, codec, resolution, unknown, length] => {
                let timestamp = timestamp.parse()?;
                let video_name = video_name.to_string();
                let codec = codec.parse()?;
                let resolution = resolution.parse()?;
                let unknown = unknown.parse()?;
                let length = length.parse()?;
                content.push(VxSeekLine {
                    timestamp,
                    video_name,
                    codec,
                    resolution,
                    unknown,
                    length,
                });
            }
            _ => return Err(ParseError::Custom(String::from("Invalid line in vxseek"))),
        }
    }
    Ok(content)
}

fn main() -> Result<(), Box<dyn Error>> {
    let args: Vec<_> = env::args().collect();
    if args.len() != 5 {
        eprintln!("Usage: {} <input.vxseek> <input.vxvideo> <resolution> <output.h264>", args[0]);
        process::exit(1);
    }
    let vxseek_filename = &args[1];
    let vxvideo_filename = &args[2];
    let wanted_resolution: Resolution = args[3].parse()?;
    let output_filename = &args[4];

    let vxseek = {
        let mut vxseek = File::open(vxseek_filename)?;
        let mut contents = String::new();
        vxseek.read_to_string(&mut contents)?;
        parse_vxseek(&contents)?
    };
    {
        let vxvideo = File::open(vxvideo_filename)?;
        let mut vxvideo = BufReader::new(vxvideo);

        let output = File::create(output_filename)?;
        let mut output = BufWriter::new(output);

        let mut name: Option<String> = None;
        for VxSeekLine { timestamp, video_name, codec, resolution, unknown, length } in vxseek {
            if let Some(ref name) = &name {
                if &video_name != name {
                    vxvideo.seek(SeekFrom::Current(length as i64))?;
                    continue;
                }
            } else {
                if resolution != wanted_resolution {
                    println!("Skipping {} due to resolution {}", video_name, resolution);
                    vxvideo.seek(SeekFrom::Current(length as i64))?;
                    continue;
                }
                name = Some(video_name);
            }
            let mut buf = Vec::with_capacity(length as usize);
            unsafe { buf.set_len(length as usize) };
            vxvideo.read_exact(&mut buf)?;
            output.write_all(&buf)?;
            output.flush()?;
        }
    }
    Ok(())
}