[nihav.git] / src / detect.rs

use std::io::SeekFrom;
use io::byteio::ByteReader;

#[derive(Debug,Clone,Copy,PartialEq)]
pub enum DetectionScore {
    No,
    ExtensionMatches,
    MagicMatches,
}

impl DetectionScore {
    pub fn less(&self, other: DetectionScore) -> bool {
        (*self as i32) < (other as i32)
    }
}

#[allow(dead_code)]
enum CondArg {
    Byte(u8),
    U16BE(u16),
    U16LE(u16),
    U24BE(u32),
    U24LE(u32),
    U32BE(u32),
    U32LE(u32),
    U64BE(u64),
    U64LE(u64),
}

impl CondArg {
    fn val(&self) -> u64 {
        match *self {
            CondArg::Byte(b) => { b as u64 }
            CondArg::U16BE(v) => { v as u64 }
            CondArg::U16LE(v) => { v as u64 }
            CondArg::U24BE(v) => { v as u64 }
            CondArg::U24LE(v) => { v as u64 }
            CondArg::U32BE(v) => { v as u64 }
            CondArg::U32LE(v) => { v as u64 }
            CondArg::U64BE(v) => { v }
            CondArg::U64LE(v) => { v }
        }
    }
    fn read_val(&self, src: &mut ByteReader) -> Option<u64> {
        match *self {
            CondArg::Byte(_) => {
                let res = src.peek_byte();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U16BE(_) => {
                let res = src.peek_u16be();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U16LE(_) => {
                let res = src.peek_u16le();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U24BE(_) => {
                let res = src.peek_u24be();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U24LE(_) => {
                let res = src.peek_u24le();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U32BE(_) => {
                let res = src.peek_u32be();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U32LE(_) => {
                let res = src.peek_u32le();
                if let Err(_) = res { return None; }
                Some(res.unwrap() as u64)
            }
            CondArg::U64BE(_) => {
                let res = src.peek_u64be();
                if let Err(_) = res { return None; }
                Some(res.unwrap())
            }
            CondArg::U64LE(_) => {
                let res = src.peek_u64le();
                if let Err(_) = res { return None; }
                Some(res.unwrap())
            }
        }
    }
    fn eq(&self, src: &mut ByteReader) -> bool {
        let val = self.read_val(src);
        if let None = val { false }
        else { val.unwrap() == self.val() }
    }
    fn ge(&self, src: &mut ByteReader) -> bool {
        let val = self.read_val(src);
        if let None = val { false }
        else { val.unwrap() >= self.val() }
    }
    fn gt(&self, src: &mut ByteReader) -> bool {
        let val = self.read_val(src);
        if let None = val { false }
        else { val.unwrap() > self.val() }
    }
    fn le(&self, src: &mut ByteReader) -> bool {
        let val = self.read_val(src);
        if let None = val { false }
        else { val.unwrap() <= self.val() }
    }
    fn lt(&self, src: &mut ByteReader) -> bool {
        let val = self.read_val(src);
        if let None = val { false }
        else { val.unwrap() < self.val() }
    }
}

#[allow(dead_code)]
enum Cond<'a> {
    And(&'a Cond<'a>, &'a Cond<'a>),
    Or(&'a Cond<'a>, &'a Cond<'a>),
    Equals(CondArg),
    EqualsString(&'static [u8]),
    InRange(CondArg, CondArg),
    LessThan(CondArg),
    GreaterThan(CondArg),
}

impl<'a> Cond<'a> {
    fn eval(&self, src: &mut ByteReader) -> bool {
        match *self {
            Cond::And(ref a, ref b) => { a.eval(src) && b.eval(src) },
            Cond::Or (ref a, ref b) => { a.eval(src) || b.eval(src) },
            Cond::Equals(ref arg)       => { arg.eq(src) },
            Cond::InRange(ref a, ref b) => { a.le(src) && b.ge(src) },
            Cond::LessThan(ref arg)     => { arg.lt(src) },
            Cond::GreaterThan(ref arg)  => { arg.gt(src) },
            Cond::EqualsString(str) => {
                let mut val: Vec<u8> = Vec::with_capacity(str.len());
                val.resize(str.len(), 0);
                let res = src.peek_buf(val.as_mut_slice());
                if let Err(_) = res { return false; }
                val == str
            }
        }
    }
}

struct CheckItem<'a> {
    offs: u32,
    cond: &'a Cond<'a>,
}

#[allow(dead_code)]
struct DetectConditions<'a> {
    demux_name: &'static str,
    extensions: &'static str,
    conditions: &'a [CheckItem<'a>],
}

const DETECTORS: &[DetectConditions] = &[
    DetectConditions {
        demux_name: "avi",
        extensions: ".avi",
        conditions: &[CheckItem{offs:0, cond: &Cond::Or(
                                                &Cond::EqualsString(b"RIFX"),
                                                &Cond::EqualsString(b"RIFF")) },
                      CheckItem{offs:8, cond: &Cond::EqualsString(b"AVI LIST") }
                     ],
    },
    DetectConditions {
        demux_name: "gdv",
        extensions: ".gdv",
        conditions: &[CheckItem{offs:0, cond: &Cond::Equals(CondArg::U32LE(0x29111994))}],
    },
];

pub fn detect_format(name: &str, src: &mut ByteReader) -> Option<(&'static str, DetectionScore)> {
    let mut result = None;
    let lname = name.to_lowercase();
    for detector in DETECTORS {
        let mut score = DetectionScore::No;
        if name.len() > 0 {
            for ext in detector.extensions.split(',') {
                if lname.ends_with(ext) {
                    score = DetectionScore::ExtensionMatches;
                    break;
                }
            }
        }
        let mut passed = detector.conditions.len() > 0;
        for ck in detector.conditions {
            let ret = src.seek(SeekFrom::Start(ck.offs as u64));
            if let Err(_) = ret {
                passed = false;
                break;
            }
            if !ck.cond.eval(src) {
                passed = false;
                break;
            }
        }
        if passed {
            score = DetectionScore::MagicMatches;
        }
        if score == DetectionScore::MagicMatches {
            return Some((detector.demux_name, score));
        }
        if let None = result {
            result = Some((detector.demux_name, score));
        } else {
            let (_, oldsc) = result.unwrap();
            if oldsc.less(score) {
                result = Some((detector.demux_name, score));
            }
        }
    }
    result
}

#[cfg(test)]
mod test {
    use super::*;
    use std::fs::File;
    use io::byteio::*;

    #[test]
    fn test_avi_detect() {
        let name = "assets/laser05.avi";
        let mut file = File::open(name).unwrap();
        let mut fr = FileReader::new_read(&mut file);
        let mut br = ByteReader::new(&mut fr);
        let (name, score) = detect_format(name, &mut br).unwrap();
        assert_eq!(name, "avi");
        assert_eq!(score, DetectionScore::MagicMatches);
    }

    #[test]
    fn test_gdv_detect() {
        let name = "assets/intro1.gdv";
        let mut file = File::open(name).unwrap();
        let mut fr = FileReader::new_read(&mut file);
        let mut br = ByteReader::new(&mut fr);
        let (name, score) = detect_format(name, &mut br).unwrap();
        assert_eq!(name, "gdv");
        assert_eq!(score, DetectionScore::MagicMatches);
    }
}
Commit	Line	Data
d8ce0de0 KS	1	use std::io::SeekFrom;
	2	use io::byteio::ByteReader;
	3
	4	#[derive(Debug,Clone,Copy,PartialEq)]
	5	pub enum DetectionScore {
	6	No,
	7	ExtensionMatches,
	8	MagicMatches,
	9	}
	10
	11	impl DetectionScore {
	12	pub fn less(&self, other: DetectionScore) -> bool {
	13	(*self as i32) < (other as i32)
	14	}
	15	}
	16
	17	#[allow(dead_code)]
	18	enum CondArg {
	19	Byte(u8),
	20	U16BE(u16),
	21	U16LE(u16),
	22	U24BE(u32),
	23	U24LE(u32),
	24	U32BE(u32),
	25	U32LE(u32),
	26	U64BE(u64),
	27	U64LE(u64),
	28	}
	29
	30	impl CondArg {
	31	fn val(&self) -> u64 {
	32	match *self {
	33	CondArg::Byte(b) => { b as u64 }
	34	CondArg::U16BE(v) => { v as u64 }
	35	CondArg::U16LE(v) => { v as u64 }
	36	CondArg::U24BE(v) => { v as u64 }
	37	CondArg::U24LE(v) => { v as u64 }
	38	CondArg::U32BE(v) => { v as u64 }
	39	CondArg::U32LE(v) => { v as u64 }
	40	CondArg::U64BE(v) => { v }
	41	CondArg::U64LE(v) => { v }
	42	}
	43	}
	44	fn read_val(&self, src: &mut ByteReader) -> Option<u64> {
	45	match *self {
	46	CondArg::Byte(_) => {
	47	let res = src.peek_byte();
	48	if let Err(_) = res { return None; }
	49	Some(res.unwrap() as u64)
	50	}
	51	CondArg::U16BE(_) => {
	52	let res = src.peek_u16be();
	53	if let Err(_) = res { return None; }
	54	Some(res.unwrap() as u64)
	55	}
	56	CondArg::U16LE(_) => {
	57	let res = src.peek_u16le();
	58	if let Err(_) = res { return None; }
	59	Some(res.unwrap() as u64)
	60	}
	61	CondArg::U24BE(_) => {
	62	let res = src.peek_u24be();
	63	if let Err(_) = res { return None; }
	64	Some(res.unwrap() as u64)
65	}
66	CondArg::U24LE(_) => {
67	let res = src.peek_u24le();
68	if let Err(_) = res { return None; }
69	Some(res.unwrap() as u64)
70	}
71	CondArg::U32BE(_) => {
72	let res = src.peek_u32be();
73	if let Err(_) = res { return None; }
74	Some(res.unwrap() as u64)
75	}
76	CondArg::U32LE(_) => {
77	let res = src.peek_u32le();
78	if let Err(_) = res { return None; }
79	Some(res.unwrap() as u64)
80	}
81	CondArg::U64BE(_) => {
82	let res = src.peek_u64be();
83	if let Err(_) = res { return None; }
84	Some(res.unwrap())
85	}
86	CondArg::U64LE(_) => {
87	let res = src.peek_u64le();
88	if let Err(_) = res { return None; }
89	Some(res.unwrap())
90	}
91	}
92	}
93	fn eq(&self, src: &mut ByteReader) -> bool {
94	let val = self.read_val(src);
95	if let None = val { false }
96	else { val.unwrap() == self.val() }
97	}
98	fn ge(&self, src: &mut ByteReader) -> bool {
99	let val = self.read_val(src);
100	if let None = val { false }
101	else { val.unwrap() >= self.val() }
102	}
103	fn gt(&self, src: &mut ByteReader) -> bool {
104	let val = self.read_val(src);
105	if let None = val { false }
106	else { val.unwrap() > self.val() }
107	}
108	fn le(&self, src: &mut ByteReader) -> bool {
109	let val = self.read_val(src);
110	if let None = val { false }
111	else { val.unwrap() <= self.val() }
112	}
113	fn lt(&self, src: &mut ByteReader) -> bool {
114	let val = self.read_val(src);
115	if let None = val { false }
116	else { val.unwrap() < self.val() }
117	}
118	}
119
120	#[allow(dead_code)]
121	enum Cond<'a> {
122	And(&'a Cond<'a>, &'a Cond<'a>),
123	Or(&'a Cond<'a>, &'a Cond<'a>),
124	Equals(CondArg),
125	EqualsString(&'static [u8]),
126	InRange(CondArg, CondArg),
127	LessThan(CondArg),
128	GreaterThan(CondArg),
129	}
130
131	impl<'a> Cond<'a> {
132	fn eval(&self, src: &mut ByteReader) -> bool {
133	match *self {
134	Cond::And(ref a, ref b) => { a.eval(src) && b.eval(src) },
135	Cond::Or (ref a, ref b) => { a.eval(src) \|\| b.eval(src) },
136	Cond::Equals(ref arg) => { arg.eq(src) },
137	Cond::InRange(ref a, ref b) => { a.le(src) && b.ge(src) },
138	Cond::LessThan(ref arg) => { arg.lt(src) },
139	Cond::GreaterThan(ref arg) => { arg.gt(src) },
140	Cond::EqualsString(str) => {
141	let mut val: Vec<u8> = Vec::with_capacity(str.len());
142	val.resize(str.len(), 0);
143	let res = src.peek_buf(val.as_mut_slice());
144	if let Err(_) = res { return false; }
145	val == str
146	}
147	}
148	}
149	}
150
151	struct CheckItem<'a> {
152	offs: u32,
153	cond: &'a Cond<'a>,
154	}
155
156	#[allow(dead_code)]
157	struct DetectConditions<'a> {
158	demux_name: &'static str,
159	extensions: &'static str,
160	conditions: &'a [CheckItem<'a>],
161	}
162
163	const DETECTORS: &[DetectConditions] = &[
164	DetectConditions {
165	demux_name: "avi",
166	extensions: ".avi",
167	conditions: &[CheckItem{offs:0, cond: &Cond::Or(
168	&Cond::EqualsString(b"RIFX"),
169	&Cond::EqualsString(b"RIFF")) },
170	CheckItem{offs:8, cond: &Cond::EqualsString(b"AVI LIST") }
171	],
172	},
173	DetectConditions {
174	demux_name: "gdv",
175	extensions: ".gdv",
176	conditions: &[CheckItem{offs:0, cond: &Cond::Equals(CondArg::U32LE(0x29111994))}],
177	},
178	];
179
180	pub fn detect_format(name: &str, src: &mut ByteReader) -> Option<(&'static str, DetectionScore)> {
181	let mut result = None;
182	let lname = name.to_lowercase();
183	for detector in DETECTORS {
184	let mut score = DetectionScore::No;
185	if name.len() > 0 {
186	for ext in detector.extensions.split(',') {
187	if lname.ends_with(ext) {
188	score = DetectionScore::ExtensionMatches;
189	break;
190	}
191	}
192	}
193	let mut passed = detector.conditions.len() > 0;
194	for ck in detector.conditions {
195	let ret = src.seek(SeekFrom::Start(ck.offs as u64));
196	if let Err(_) = ret {
197	passed = false;
198	break;
199	}
200	if !ck.cond.eval(src) {
201	passed = false;
202	break;
203	}
204	}
205	if passed {
206	score = DetectionScore::MagicMatches;
207	}
208	if score == DetectionScore::MagicMatches {
209	return Some((detector.demux_name, score));
210	}
211	if let None = result {
212	result = Some((detector.demux_name, score));
213	} else {
214	let (_, oldsc) = result.unwrap();
215	if oldsc.less(score) {
216	result = Some((detector.demux_name, score));
217	}
218	}
219	}
220	result
221	}
222
223	#[cfg(test)]
224	mod test {
225	use super::*;
226	use std::fs::File;
227	use io::byteio::*;
228
229	#[test]
230	fn test_avi_detect() {
231	let name = "assets/laser05.avi";
232	let mut file = File::open(name).unwrap();
233	let mut fr = FileReader::new_read(&mut file);
234	let mut br = ByteReader::new(&mut fr);
235	let (name, score) = detect_format(name, &mut br).unwrap();
236	assert_eq!(name, "avi");
237	assert_eq!(score, DetectionScore::MagicMatches);
238	}
239
240	#[test]
241	fn test_gdv_detect() {
242	let name = "assets/intro1.gdv";
243	let mut file = File::open(name).unwrap();
244	let mut fr = FileReader::new_read(&mut file);
245	let mut br = ByteReader::new(&mut fr);
246	let (name, score) = detect_format(name, &mut br).unwrap();
247	assert_eq!(name, "gdv");
248	assert_eq!(score, DetectionScore::MagicMatches);
249	}
250	}