nihav-core/src/detect.rs

   1 use std::io::SeekFrom;
   2 use crate::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 Arg {
  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 Arg {
  31     fn val(&self) -> u64 {
  32         match *self {
  33             Arg::Byte(b) => { u64::from(b) }
  34             Arg::U16BE(v) => { u64::from(v) }
  35             Arg::U16LE(v) => { u64::from(v) }
  36             Arg::U24BE(v) => { u64::from(v) }
  37             Arg::U24LE(v) => { u64::from(v) }
  38             Arg::U32BE(v) => { u64::from(v) }
  39             Arg::U32LE(v) => { u64::from(v) }
  40             Arg::U64BE(v) => { v }
  41             Arg::U64LE(v) => { v }
  42         }
  43     }
  44     fn read_val(&self, src: &mut ByteReader) -> Option<u64> {
  45         match *self {
  46             Arg::Byte(_) => {
  47                 let res = src.peek_byte();
  48                 if res.is_err() { return None; }
  49                 Some(u64::from(res.unwrap()))
  50             }
  51             Arg::U16BE(_) => {
  52                 let res = src.peek_u16be();
  53                 if res.is_err() { return None; }
  54                 Some(u64::from(res.unwrap()))
  55             }
  56             Arg::U16LE(_) => {
  57                 let res = src.peek_u16le();
  58                 if res.is_err() { return None; }
  59                 Some(u64::from(res.unwrap()))
  60             }
  61             Arg::U24BE(_) => {
  62                 let res = src.peek_u24be();
  63                 if res.is_err() { return None; }
  64                 Some(u64::from(res.unwrap()))
  65             }
  66             Arg::U24LE(_) => {
  67                 let res = src.peek_u24le();
  68                 if res.is_err() { return None; }
  69                 Some(u64::from(res.unwrap()))
  70             }
  71             Arg::U32BE(_) => {
  72                 let res = src.peek_u32be();
  73                 if res.is_err() { return None; }
  74                 Some(u64::from(res.unwrap()))
  75             }
  76             Arg::U32LE(_) => {
  77                 let res = src.peek_u32le();
  78                 if res.is_err() { return None; }
  79                 Some(u64::from(res.unwrap()))
  80             }
  81             Arg::U64BE(_) => {
  82                 let res = src.peek_u64be();
  83                 if res.is_err() { return None; }
  84                 Some(res.unwrap())
  85             }
  86             Arg::U64LE(_) => {
  87                 let res = src.peek_u64le();
  88                 if res.is_err() { 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 val.is_none() { 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 val.is_none() { 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 val.is_none() { 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 val.is_none() { 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 val.is_none() { false }
 116         else { val.unwrap() < self.val() }
 117     }
 118 }
 119
 120 #[allow(dead_code)]
 121 enum CC<'a> {
 122     Or(&'a CC<'a>, &'a CC<'a>),
 123     Eq(Arg),
 124     Str(&'static [u8]),
 125     In(Arg, Arg),
 126     Lt(Arg),
 127     Le(Arg),
 128     Gt(Arg),
 129     Ge(Arg),
 130 }
 131
 132 impl<'a> CC<'a> {
 133     fn eval(&self, src: &mut ByteReader) -> bool {
 134         match *self {
 135             CC::Or (ref a, ref b) => { a.eval(src) || b.eval(src) },
 136             CC::Eq(ref arg)      => { arg.eq(src) },
 137             CC::In(ref a, ref b) => { a.ge(src) && b.le(src) },
 138             CC::Lt(ref arg)      => { arg.lt(src) },
 139             CC::Le(ref arg)      => { arg.le(src) },
 140             CC::Gt(ref arg)      => { arg.gt(src) },
 141             CC::Ge(ref arg)      => { arg.ge(src) },
 142             CC::Str(str) => {
 143                 let mut val: Vec<u8> = vec![0; str.len()];
 144                 let res = src.peek_buf(val.as_mut_slice());
 145                 if res.is_err() { return false; }
 146                 val == str
 147             }
 148         }
 149     }
 150 }
 151
 152 struct CheckItem<'a> {
 153     offs: u32,
 154     cond: &'a CC<'a>,
 155 }
 156
 157 #[allow(dead_code)]
 158 struct DetectConditions<'a> {
 159     demux_name: &'static str,
 160     extensions: &'static str,
 161     conditions: &'a [CheckItem<'a>],
 162 }
 163
 164 const DETECTORS: &[DetectConditions] = &[
 165     DetectConditions {
 166         demux_name: "avi",
 167         extensions: ".avi",
 168         conditions: &[CheckItem{offs: 0, cond: &CC::Or(&CC::Str(b"RIFF"), &CC::Str(b"ON2 ")) },
 169                       CheckItem{offs: 8, cond: &CC::Or(&CC::Or(&CC::Str(b"AVI LIST"),
 170                                                                &CC::Str(b"AVIXLIST")),
 171                                                                &CC::Str(b"ON2fLIST")) },
 172                      ]
 173     },
 174     DetectConditions {
 175         demux_name: "gdv",
 176         extensions: ".gdv",
 177         conditions: &[CheckItem{offs: 0, cond: &CC::Eq(Arg::U32LE(0x29111994))}],
 178     },
 179     DetectConditions {
 180         demux_name: "realaudio",
 181         extensions: ".ra,.ram",
 182         conditions: &[CheckItem{offs: 0, cond: &CC::Str(b".ra\xFD")}],
 183     },
 184     DetectConditions {
 185         demux_name: "realmedia",
 186         extensions: ".rm,.rmvb,.rma,.ra,.ram",
 187         conditions: &[CheckItem{offs: 0, cond: &CC::Or(&CC::Str(b".RMF"), &CC::Str(b".RMP")) },
 188                       CheckItem{offs: 4, cond: &CC::Ge(Arg::U32BE(10))}],
 189     },
 190     DetectConditions {
 191         demux_name: "real_ivr",
 192         extensions: ".ivr",
 193         conditions: &[CheckItem{offs: 0, cond: &CC::Or(&CC::Str(b".R1M"), &CC::Str(b".REC"))}],
 194     },
 195     DetectConditions {
 196         demux_name: "bink",
 197         extensions: ".bik,.bk2",
 198         conditions: &[CheckItem{offs: 0, cond: &CC::Or(&CC::In(Arg::U32BE(0x32494B62),     // BIKb
 199                                                                Arg::U32BE(0x32494B7B)),    // BIKz
 200                                                        &CC::In(Arg::U32BE(0x4B423261),     // KB2a
 201                                                                Arg::U32BE(0x4B42327B)))}], // KB2z
 202     },
 203     DetectConditions {
 204         demux_name: "smacker",
 205         extensions: ".smk",
 206         conditions: &[CheckItem{offs: 0, cond: &CC::Or(&CC::Str(b"SMK2"), &CC::Str(b"SMK4"))}],
 207     },
 208     DetectConditions {
 209         demux_name: "bmv",
 210         extensions: ".bmv",
 211         conditions: &[],
 212     },
 213     DetectConditions {
 214         demux_name: "bmv3",
 215         extensions: ".bmv",
 216         conditions: &[CheckItem{offs: 0, cond: &CC::Str(b"BMVi") },
 217                       CheckItem{offs: 32, cond: &CC::Str(b"DATA")}],
 218     },
 219     DetectConditions {
 220         demux_name: "vmd",
 221         extensions: ".vmd",
 222         conditions: &[],
 223     },
 224 ];
 225
 226 pub fn detect_format(name: &str, src: &mut ByteReader) -> Option<(&'static str, DetectionScore)> {
 227     let mut result = None;
 228     let lname = name.to_lowercase();
 229     for detector in DETECTORS {
 230         let mut score = DetectionScore::No;
 231         if !name.is_empty() {
 232             for ext in detector.extensions.split(',') {
 233                 if lname.ends_with(ext) {
 234                     score = DetectionScore::ExtensionMatches;
 235                     break;
 236                 }
 237             }
 238         }
 239         let mut passed = !detector.conditions.is_empty();
 240         for ck in detector.conditions {
 241             let ret = src.seek(SeekFrom::Start(u64::from(ck.offs)));
 242             if ret.is_err() {
 243                 passed = false;
 244                 break;
 245             }
 246             if !ck.cond.eval(src) {
 247                 passed = false;
 248                 break;
 249             }
 250         }
 251         if passed {
 252             score = DetectionScore::MagicMatches;
 253         }
 254         if score == DetectionScore::MagicMatches {
 255             return Some((detector.demux_name, score));
 256         }
 257         if result.is_none() && score != DetectionScore::No {
 258             result = Some((detector.demux_name, score));
 259         } else if result.is_some() {
 260             let (_, oldsc) = result.unwrap();
 261             if oldsc.less(score) {
 262                 result = Some((detector.demux_name, score));
 263             }
 264         }
 265     }
 266     result
 267 }
 268
 269 #[cfg(test)]
 270 mod test {
 271     use super::*;
 272     use std::fs::File;
 273     use crate::io::byteio::*;
 274
 275     #[test]
 276     fn test_avi_detect() {
 277         let name = "assets/Indeo/laser05.avi";
 278         let mut file = File::open(name).unwrap();
 279         let mut fr = FileReader::new_read(&mut file);
 280         let mut br = ByteReader::new(&mut fr);
 281         let (name, score) = detect_format(name, &mut br).unwrap();
 282         assert_eq!(name, "avi");
 283         assert_eq!(score, DetectionScore::MagicMatches);
 284     }
 285
 286     #[test]
 287     fn test_gdv_detect() {
 288         let name = "assets/Game/intro1.gdv";
 289         let mut file = File::open(name).unwrap();
 290         let mut fr = FileReader::new_read(&mut file);
 291         let mut br = ByteReader::new(&mut fr);
 292         let (name, score) = detect_format(name, &mut br).unwrap();
 293         assert_eq!(name, "gdv");
 294         assert_eq!(score, DetectionScore::MagicMatches);
 295     }
 296 }