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