flamberge_schemes/
lib.rs

1//! Per-scheme DRM removal. Each module combines a parsed container
2//! (`flamberge-formats`), candidate keys (`flamberge-keys`), and ciphers
3//! (`flamberge-crypto`) to produce a decrypted book.
4//!
5//! Reference: `docs/DEDRM_SCHEMES.md`.
6
7pub mod adept;
8mod epub_common;
9pub mod ereader;
10pub mod error;
11pub mod ignoble;
12pub mod kfx;
13pub mod kobo;
14pub mod mobipocket;
15mod pdf_common;
16pub mod topaz;
17
18pub use error::SchemeError;
19pub use flamberge_keys::KeyStore;
20
21pub type Result<T> = std::result::Result<T, SchemeError>;
22
23/// A decrypted book ready to be written out.
24#[derive(Debug, Clone)]
25pub struct DecryptedBook {
26    pub data: Vec<u8>,
27    /// Output file extension without the dot (e.g. `mobi`, `epub`, `pdf`).
28    pub extension: String,
29    /// The book's display title, when the scheme can recover it. Used by the
30    /// CLI for title-based output naming; `None` when unavailable.
31    pub title: Option<String>,
32}
33
34/// The DRM schemes this tool can (eventually) remove.
35#[derive(Debug, Clone, Copy, PartialEq, Eq)]
36pub enum Scheme {
37    Mobipocket,
38    Topaz,
39    Kfx,
40    AdeptEpub,
41    IgnobleEpub,
42    AdeptPdf,
43    IgnoblePdf,
44    EReader,
45    Kobo,
46}
47
48/// Candidate schemes to try for a given file extension, in priority order.
49/// Kindle-family files disambiguate further by magic inside the handler.
50///
51/// Mirrors `DeDRM_plugin/__init__.py::run` extension routing.
52pub fn candidates_for_extension(ext: &str) -> &'static [Scheme] {
53    match ext.trim_start_matches('.').to_ascii_lowercase().as_str() {
54        "prc" | "mobi" | "pobi" | "azw" | "azw1" | "azw3" | "azw4" | "tpz" | "kfx-zip" => {
55            &[Scheme::Mobipocket, Scheme::Topaz, Scheme::Kfx]
56        }
57        "pdb" => &[Scheme::EReader],
58        "pdf" => &[Scheme::IgnoblePdf, Scheme::AdeptPdf],
59        "epub" => &[Scheme::IgnobleEpub, Scheme::AdeptEpub],
60        "kepub" => &[Scheme::Kobo],
61        _ => &[],
62    }
63}
64
65/// Refine a Kindle-family choice by container magic (see
66/// `k4mobidedrm.py::GetDecryptedBook`).
67pub fn kindle_scheme_from_magic(data: &[u8]) -> Option<Scheme> {
68    if data.starts_with(b"\xeaDRMION\xee") {
69        // A bare DRMION cannot be decrypted on its own.
70        None
71    } else if data.starts_with(b"PK\x03\x04") {
72        Some(Scheme::Kfx)
73    } else if data.starts_with(b"TPZ") {
74        Some(Scheme::Topaz)
75    } else {
76        Some(Scheme::Mobipocket)
77    }
78}
79
80/// Decrypt `input` given its file extension and a key store, trying each
81/// candidate scheme until one succeeds.
82pub fn decrypt(input: &[u8], ext: &str, keys: &KeyStore) -> Result<DecryptedBook> {
83    let mut candidates = candidates_for_extension(ext).to_vec();
84
85    // Real Kobo books are named `*.kepub.epub` (extension `epub`) or an
86    // extension-less volume id, so they would otherwise route to the EPUB
87    // schemes — or nowhere — and never reach the Kobo handler. A supplied Kobo
88    // DB is the unambiguous "this is a Kobo book" signal, so for a ZIP input we
89    // append Kobo as a fallback candidate (tried after any EPUB schemes, and not
90    // for Kindle-family archives like `.kfx-zip`, which have their own routing).
91    let is_zip = input.starts_with(b"PK\x03\x04");
92    let is_kindle_family = candidates
93        .iter()
94        .any(|s| matches!(s, Scheme::Mobipocket | Scheme::Topaz | Scheme::Kfx));
95    if keys.kobo_db.is_some() && is_zip && !is_kindle_family && !candidates.contains(&Scheme::Kobo)
96    {
97        candidates.push(Scheme::Kobo);
98    }
99
100    if candidates.is_empty() {
101        return Err(SchemeError::UnknownFormat(ext.to_string()));
102    }
103
104    for scheme in candidates {
105        let result = match scheme {
106            Scheme::Mobipocket => mobipocket::decrypt(input, keys),
107            Scheme::Topaz => topaz::decrypt(input, keys),
108            Scheme::Kfx => kfx::decrypt(input, keys),
109            Scheme::AdeptEpub => adept::decrypt_epub(input, keys),
110            Scheme::IgnobleEpub => ignoble::decrypt_epub(input, keys),
111            Scheme::AdeptPdf => adept::decrypt_pdf(input, keys),
112            Scheme::IgnoblePdf => ignoble::decrypt_pdf(input, keys),
113            Scheme::EReader => ereader::decrypt(input, keys),
114            Scheme::Kobo => kobo::decrypt(input, keys),
115        };
116        // `NotThisScheme` is the only "keep looking" signal: it means the file
117        // isn't handled by this scheme. Any other error means the scheme claimed
118        // the file (right magic/structure) but decryption failed — that verdict
119        // is terminal and must not be masked by a later candidate's error.
120        match result {
121            Ok(book) => return Ok(book),
122            Err(SchemeError::NotThisScheme) => continue,
123            Err(e) => return Err(e),
124        }
125    }
126    // No candidate recognized the file.
127    Err(SchemeError::NoKeyWorked)
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133
134    #[test]
135    fn extension_routing_matches_plugin() {
136        assert_eq!(candidates_for_extension("mobi").len(), 3);
137        assert_eq!(candidates_for_extension(".AZW3").len(), 3); // dotted + upper-case
138        assert_eq!(candidates_for_extension("pdb"), &[Scheme::EReader]);
139        assert_eq!(
140            candidates_for_extension("pdf"),
141            &[Scheme::IgnoblePdf, Scheme::AdeptPdf]
142        );
143        assert_eq!(
144            candidates_for_extension("epub"),
145            &[Scheme::IgnobleEpub, Scheme::AdeptEpub]
146        );
147        assert_eq!(candidates_for_extension("kepub"), &[Scheme::Kobo]);
148        assert!(candidates_for_extension("txt").is_empty());
149    }
150
151    #[test]
152    fn unknown_extension_is_reported() {
153        let keys = KeyStore::new();
154        match decrypt(b"whatever", "txt", &keys) {
155            Err(SchemeError::UnknownFormat(ext)) => assert_eq!(ext, "txt"),
156            other => panic!("expected UnknownFormat, got {other:?}"),
157        }
158    }
159
160    #[test]
161    fn all_candidates_falling_through_yields_no_key_worked() {
162        // A non-ZIP buffer is not an EPUB, so both EPUB candidates return
163        // `NotThisScheme`; dispatch must exhaust them and report `NoKeyWorked`
164        // rather than surfacing an internal error from the last candidate tried.
165        let keys = KeyStore::new();
166        assert!(matches!(
167            decrypt(b"definitely not a zip", "epub", &keys),
168            Err(SchemeError::NoKeyWorked)
169        ));
170    }
171
172    #[test]
173    fn terminal_error_is_surfaced_not_masked() {
174        // A scheme that claims the file but fails must have its error propagated
175        // verbatim (not collapsed to `NoKeyWorked`). `.pdb` routes to EReader,
176        // whose first step is `PalmDb::parse(input)?`; a truncated buffer makes
177        // that a terminal `Format` error, which dispatch must surface as-is.
178        let keys = KeyStore::new();
179        assert!(matches!(
180            decrypt(&[0u8; 8], "pdb", &keys),
181            Err(SchemeError::Format(_))
182        ));
183    }
184}
flamberge_schemes/lib.rs

flamberge_schemes/
lib.rs
