use crate::error::{DocError, Result};
const OLE_MAGIC: &[u8] = b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1";
const ZIP_MAGIC: &[u8] = b"PK\x03\x04";
const PDF_MAGIC: &[u8] = b"%PDF";
const OOXML: &[&str] = &["docx", "xlsx", "pptx"];
const ODF: &[&str] = &["odt", "ods", "odp"];
const OLE_EXT: &[&str] = &["doc", "ppt", "xls"];
const TEXT_EXT: &[&str] = &["txt", "md", "markdown", "csv", "html", "htm"];
const SUB_EXT: &[&str] = &["srt", "vtt", "ass", "ssa"];
const RTF_EXT: &[&str] = &["rtf"];
fn head_of(path: &str, n: usize) -> Result<Vec<u8>> {
std::fs::read(path)
.map(|v| v[..v.len().min(n)].to_vec())
.map_err(|e| DocError::corrupt(format!("cannot read {path}: {e}")))
}
pub fn sniff(path: &str) -> Result<&'static str> {
let ext = path
.rsplit('.')
.next()
.map(|e| e.to_ascii_lowercase())
.unwrap_or_default();
let h = head_of(path, 8)?;
if h.starts_with(OLE_MAGIC) {
if OOXML.contains(&ext.as_str()) {
return Err(DocError::encrypted(format!("{ext} 为加密容器(OLE)")));
}
if OLE_EXT.contains(&ext.as_str()) {
return Ok("ole");
}
return Err(DocError::unsupported(format!(
"OLE 容器但扩展名不支持: {ext}"
)));
}
if h.starts_with(ZIP_MAGIC) {
if ext == "epub" {
return Ok("epub");
}
if ODF.contains(&ext.as_str()) {
return Ok(match ext.as_str() {
"odt" => "odt",
"ods" => "ods",
_ => "odp",
});
}
if OOXML.contains(&ext.as_str()) {
return Ok(match ext.as_str() {
"docx" => "docx",
"xlsx" => "xlsx",
_ => "pptx",
});
}
return Err(DocError::corrupt("ZIP 容器但缺少 OOXML 关键部件"));
}
if h.starts_with(PDF_MAGIC) {
return Ok("pdf");
}
if RTF_EXT.contains(&ext.as_str()) {
return Ok("rtf");
}
if OLE_EXT.contains(&ext.as_str()) {
return Err(DocError::unsupported(format!(
"{ext} 文件头不是 OLE,疑似伪装"
)));
}
if TEXT_EXT.contains(&ext.as_str()) {
let kind: &'static str = match ext.as_str() {
"txt" => "txt",
"md" | "markdown" => "md",
"csv" => "csv",
"html" | "htm" => "html",
_ => "txt",
};
return Ok(kind);
}
if SUB_EXT.contains(&ext.as_str()) {
return Ok(subtitle_format(&ext));
}
Ok("txt") }
fn subtitle_format(ext: &str) -> &'static str {
match ext {
"vtt" => "vtt",
"ass" => "ass",
"ssa" => "ssa",
_ => "srt",
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
#[test]
fn sniff_pdf_and_ole() {
let d = std::env::temp_dir();
let pdf = d.join("docparser_router_test.pdf");
let mut f = std::fs::File::create(&pdf).unwrap();
f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
drop(f);
assert_eq!(sniff(pdf.to_str().unwrap()).unwrap(), "pdf");
std::fs::remove_file(&pdf).ok();
let renamed = d.join("renamed_pdf.doc");
let mut f = std::fs::File::create(&renamed).unwrap();
f.write_all(b"%PDF-1.7\n%%EOF").unwrap();
drop(f);
assert_eq!(sniff(renamed.to_str().unwrap()).unwrap(), "pdf");
std::fs::remove_file(&renamed).ok();
let fake = d.join("fake.doc");
let mut f = std::fs::File::create(&fake).unwrap();
f.write_all(b"NOTANOFFICEFILE").unwrap();
drop(f);
let err = sniff(fake.to_str().unwrap()).unwrap_err();
assert_eq!(err.code.as_str(), "unsupported");
std::fs::remove_file(&fake).ok();
}
#[test]
fn ooxml_ole_container_is_encrypted() {
let d = std::env::temp_dir();
let fx = d.join("enc.docx");
let mut f = std::fs::File::create(&fx).unwrap();
f.write_all(OLE_MAGIC).unwrap();
drop(f);
let err = sniff(fx.to_str().unwrap()).unwrap_err();
assert_eq!(err.code.as_str(), "encrypted");
std::fs::remove_file(&fx).ok();
}
#[test]
fn subtitle_extensions_route_to_subtitle_formats() {
let d = std::env::temp_dir();
for (name, want) in [
("r.srt", "srt"),
("r.vtt", "vtt"),
("r.ass", "ass"),
("r.ssa", "ssa"),
] {
let p = d.join(format!("docparser_router_{name}"));
let mut f = std::fs::File::create(&p).unwrap();
f.write_all(b"1\n00:00:01,000 --> 00:00:02,000\nhi\n")
.unwrap();
drop(f);
assert_eq!(sniff(p.to_str().unwrap()).unwrap(), want, "{name}");
std::fs::remove_file(&p).ok();
}
}
#[test]
fn rtf_extension_routes_to_rtf() {
let d = std::env::temp_dir();
let p = d.join("docparser_router_r.rtf");
std::fs::write(&p, br"{\rtf1\ansi hi\par}").unwrap();
assert_eq!(sniff(p.to_str().unwrap()).unwrap(), "rtf");
std::fs::remove_file(&p).ok();
}
#[test]
fn zip_container_routes_odf_family() {
let d = std::env::temp_dir();
for (name, want) in [
("r.odt", "odt"),
("r.ods", "ods"),
("r.odp", "odp"),
("r.docx", "docx"),
] {
let p = d.join(format!("docparser_router_{name}"));
let mut f = std::fs::File::create(&p).unwrap();
f.write_all(ZIP_MAGIC).unwrap();
drop(f);
assert_eq!(sniff(p.to_str().unwrap()).unwrap(), want, "{name}");
std::fs::remove_file(&p).ok();
}
}
}