refactor: consolidate platform engines into senbei-engine

senbei-engine/src/android/extract/pipeline.rs

@@ -0,0 +1,538 @@
+use std::collections::{BTreeMap, BTreeSet, HashSet};
+use std::fs::{File, create_dir_all};
+use std::io::Write;
+use std::path::{Path, PathBuf};
+
+use memmap2::MmapOptions;
+use senbei_crypto::android::{Module9bConfig, decode_container};
+use serde_json::to_vec_pretty;
+use sha2::{Digest, Sha256};
+use tempfile::NamedTempFile;
+
+use super::error::{Error, Result, invalid};
+use super::report::{
+    ArtifactReport, DecoderReport, ExtractionReport, ModuleRegistryEntry, RecordReport,
+    Stage1Report, StreamParent, StreamReport,
+};
+use super::stage1::{
+    DEFAULT_CIPHER_CONSTANT, DEFAULT_OUTER_SIZE, SHT_LOUSER, Stage1Result, inspect,
+};
+use super::stream::{DIRECT_FLAG, Record, parse_record_stream};
+
+/// Inputs and output locations for one complete static Stage 2 extraction.
+#[derive(Debug, Clone)]
+pub struct ExtractOptions {
+    pub input: PathBuf,
+    pub output_dir: PathBuf,
+    pub stage2_output: Option<PathBuf>,
+    pub outer_size: usize,
+    pub cipher_constant: u32,
+}
+
+impl ExtractOptions {
+    #[must_use]
+    pub fn with_defaults(input: PathBuf, output_dir: PathBuf) -> Self {
+        Self {
+            input,
+            output_dir,
+            stage2_output: None,
+            outer_size: DEFAULT_OUTER_SIZE,
+            cipher_constant: DEFAULT_CIPHER_CONSTANT,
+        }
+    }
+}
+
+#[derive(Debug)]
+struct LoadedModule {
+    image: Vec<u8>,
+    metadata: Option<Vec<u8>>,
+    image_path: String,
+    metadata_path: Option<String>,
+    sha256: String,
+    depth: usize,
+    record_index: usize,
+    command_id: u32,
+    init_offset: u32,
+    entry_offset: u32,
+}
+
+#[derive(Debug, Clone, Copy)]
+struct ArtifactSpec<'a> {
+    suffix: &'a str,
+    kind: &'a str,
+    classification: &'a str,
+}
+
+struct Extractor {
+    output_dir: PathBuf,
+    streams: Vec<StreamReport>,
+    artifacts: Vec<ArtifactReport>,
+    registry: BTreeMap<u32, LoadedModule>,
+    seen_streams: HashSet<(u32, String)>,
+}
+
+pub fn extract_stage2(options: &ExtractOptions) -> Result<ExtractionReport> {
+    let input_path = absolute(&options.input)?;
+    let output_dir = absolute(&options.output_dir)?;
+    if !input_path.is_file() {
+        return invalid(format!(
+            "protected ELF does not exist: {}",
+            input_path.display()
+        ));
+    }
+    if let Some(stage2_output) = &options.stage2_output {
+        let stage2_output = absolute(stage2_output)?;
+        if stage2_output == input_path {
+            return invalid("refusing to overwrite the protected ELF with Stage 2 output");
+        }
+    }
+    create_dir_all(&output_dir)
+        .map_err(|source| Error::io("create Stage 2 output directory", &output_dir, source))?;
+
+    let file = File::open(&input_path)
+        .map_err(|source| Error::io("open protected ELF", &input_path, source))?;
+    // SAFETY: the mapping is read-only, the file remains open for the mapping
+    // lifetime, and extraction never mutates or truncates the source.
+    let source = unsafe { MmapOptions::new().map(&file) }
+        .map_err(|source| Error::io("map protected ELF", &input_path, source))?;
+    let stage1 = inspect(
+        &source,
+        &input_path,
+        options.outer_size,
+        options.cipher_constant,
+    )?;
+    if let Some(stage2_output) = &options.stage2_output {
+        write_atomic(&absolute(stage2_output)?, &stage1.plaintext)?;
+    }
+
+    let core_config =
+        Module9bConfig::parse_embedded(&stage1.plaintext).map_err(Error::EmbeddedConfig)?;
+    let bootstrap_end = stage1
+        .remaining_file_offset
+        .checked_add(stage1.remaining_size)
+        .ok_or_else(|| Error::Invalid("Stage 2 bootstrap range overflow".to_owned()))?;
+    let bootstrap = source
+        .get(stage1.remaining_file_offset..bootstrap_end)
+        .ok_or_else(|| Error::Invalid("Stage 2 bootstrap range is outside the ELF".to_owned()))?;
+    let mut extractor = Extractor {
+        output_dir: output_dir.clone(),
+        streams: Vec::new(),
+        artifacts: Vec::new(),
+        registry: BTreeMap::new(),
+        seen_streams: HashSet::new(),
+    };
+    extractor.extract_stream(
+        bootstrap,
+        0xe2,
+        0,
+        None,
+        Some(stage1.remaining_file_offset),
+        core_config,
+    )?;
+
+    let module_registry = extractor
+        .registry
+        .values()
+        .map(|module| ModuleRegistryEntry {
+            command_id: module.command_id,
+            size: module.image.len(),
+            sha256: module.sha256.clone(),
+            depth: module.depth,
+            record_index: module.record_index,
+            image_path: module.image_path.clone(),
+            metadata_path: module.metadata_path.clone(),
+            init_offset: module.init_offset,
+            entry_offset: module.entry_offset,
+            classification: if module.metadata.is_some() {
+                "module_image".to_owned()
+            } else {
+                "decoded_data".to_owned()
+            },
+        })
+        .collect::<Vec<_>>();
+    let report = ExtractionReport {
+        format_version: 4,
+        protected_elf: input_path.display().to_string(),
+        output_dir: output_dir.display().to_string(),
+        stage1: stage1_report(&stage1, options.outer_size),
+        streams: extractor.streams,
+        artifacts: extractor.artifacts,
+        errors: Vec::new(),
+        module_registry,
+    };
+    write_json_atomic(&output_dir.join("index.json"), &report)?;
+    Ok(report)
+}
+
+impl Extractor {
+    fn extract_stream(
+        &mut self,
+        stream: &[u8],
+        stream_id: u32,
+        depth: usize,
+        parent: Option<StreamParent>,
+        source_file_offset: Option<usize>,
+        config: Module9bConfig,
+    ) -> Result<()> {
+        let digest = sha256(stream);
+        if !self.seen_streams.insert((stream_id, digest.clone())) {
+            return Ok(());
+        }
+        let (header, records, table_size) =
+            parse_record_stream(stream, stream_id).map_err(|source| {
+                Error::Invalid(format!(
+                    "depth {depth} stream 0x{stream_id:02X} record table: {source}"
+                ))
+            })?;
+        let mut stream_report = StreamReport {
+            depth,
+            stream_id,
+            parent,
+            source_file_offset,
+            available_size: stream.len(),
+            descriptor_table_size: table_size,
+            encrypted_header_words: header.encrypted_words,
+            decrypted_header_words: header.decrypted_words,
+            record_state: header.record_state,
+            sha256: digest,
+            decoder: decoder_report(
+                if depth == 0 {
+                    "embedded_stage2"
+                } else {
+                    "decoded_interpreter"
+                },
+                (depth != 0).then_some(stream_id),
+                &config,
+            ),
+            records: Vec::with_capacity(records.len()),
+        };
+        let mut direct_records = Vec::new();
+        let mut modules_at_level = BTreeSet::new();
+
+        for record in records {
+            let mut result = record_report(record);
+            let mut image_data = None;
+            let mut metadata_data = None;
+
+            if !record.direct() && record.image_size != 0 {
+                let image_source = record_tail(stream, record.image_offset)?;
+                let image = decode_container(image_source, &config, record.image_size as usize)
+                    .map_err(|source| Error::RecordDecode {
+                        depth,
+                        stream_id,
+                        record_index: record.index,
+                        command_id: record.command_id,
+                        part: "image decode",
+                        source,
+                    })?;
+                let classification = if record.metadata_size != 0 {
+                    "module_image"
+                } else {
+                    "decoded_data"
+                };
+                let artifact = self.write_artifact(
+                    &record,
+                    depth,
+                    stream_id,
+                    ArtifactSpec {
+                        suffix: "module.bin",
+                        kind: "decoded_container",
+                        classification,
+                    },
+                    &image,
+                )?;
+                result.image = Some(artifact.clone());
+                image_data = Some((image, artifact));
+            }
+            if record.metadata_size != 0 {
+                let metadata_source = record_tail(stream, record.metadata_offset)?;
+                let metadata =
+                    decode_container(metadata_source, &config, record.metadata_size as usize)
+                        .map_err(|source| Error::RecordDecode {
+                            depth,
+                            stream_id,
+                            record_index: record.index,
+                            command_id: record.command_id,
+                            part: "metadata decode",
+                            source,
+                        })?;
+                let artifact = self.write_artifact(
+                    &record,
+                    depth,
+                    stream_id,
+                    ArtifactSpec {
+                        suffix: "metadata.bin",
+                        kind: "decoded_metadata",
+                        classification: "decoded_metadata",
+                    },
+                    &metadata,
+                )?;
+                result.metadata = Some(artifact.clone());
+                metadata_data = Some((metadata, artifact));
+            }
+            if let Some((image, image_artifact)) = image_data {
+                let (metadata, metadata_path) = if let Some((data, artifact)) = metadata_data {
+                    (Some(data), Some(artifact.path))
+                } else {
+                    (None, None)
+                };
+                self.register_module(LoadedModule {
+                    sha256: image_artifact.sha256.clone(),
+                    image_path: image_artifact.path.clone(),
+                    metadata_path,
+                    image,
+                    metadata,
+                    depth,
+                    record_index: record.index,
+                    command_id: record.command_id,
+                    init_offset: record.init_offset,
+                    entry_offset: record.entry_offset,
+                })?;
+                modules_at_level.insert(record.command_id);
+            }
+            if record.direct() && record.image_size != 0 {
+                direct_records.push((record, stream_report.records.len()));
+            }
+            stream_report.records.push(result);
+        }
+
+        let mut children = Vec::new();
+        for (record, report_index) in direct_records {
+            let next_stream_id = record.command_id.wrapping_sub(0x10);
+            if modules_at_level.contains(&next_stream_id) {
+                stream_report.records[report_index].nested_stream_id = Some(next_stream_id);
+                children.push((record, next_stream_id));
+                continue;
+            }
+            let direct_data = record_slice(stream, record.image_offset, record.image_size)?;
+            let artifact = self.write_artifact(
+                &record,
+                depth,
+                stream_id,
+                ArtifactSpec {
+                    suffix: "direct.bin",
+                    kind: "direct",
+                    classification: "direct_data",
+                },
+                direct_data,
+            )?;
+            stream_report.records[report_index].image = Some(artifact);
+        }
+
+        self.streams.push(stream_report);
+        for (record, next_stream_id) in children {
+            let child_data = record_slice(stream, record.image_offset, record.image_size)?;
+            let parent = StreamParent {
+                stream_id,
+                record_index: record.index,
+                command_id: record.command_id,
+            };
+            let interpreter = self.registry.get(&next_stream_id).ok_or_else(|| {
+                Error::Invalid(format!(
+                    "depth {depth} stream 0x{stream_id:02X} child 0x{next_stream_id:02X} has no interpreter module"
+                ))
+            })?;
+            let interpreter_config =
+                Module9bConfig::parse(&interpreter.image).map_err(|source| {
+                    Error::InterpreterConfig {
+                        depth: depth + 1,
+                        stream_id: next_stream_id,
+                        interpreter_id: next_stream_id,
+                        source,
+                    }
+                })?;
+            self.extract_stream(
+                child_data,
+                next_stream_id,
+                depth + 1,
+                Some(parent),
+                None,
+                interpreter_config,
+            )?;
+        }
+        Ok(())
+    }
+
+    fn register_module(&mut self, module: LoadedModule) -> Result<()> {
+        if let Some(previous) = self.registry.get(&module.command_id) {
+            if previous.sha256 != module.sha256 {
+                return invalid(format!(
+                    "module 0x{:02X} produced conflicting images: {} and {}",
+                    module.command_id, previous.sha256, module.sha256
+                ));
+            }
+            return Ok(());
+        }
+        self.registry.insert(module.command_id, module);
+        Ok(())
+    }
+
+    fn write_artifact(
+        &mut self,
+        record: &Record,
+        depth: usize,
+        stream_id: u32,
+        spec: ArtifactSpec<'_>,
+        data: &[u8],
+    ) -> Result<ArtifactReport> {
+        let digest = sha256(data);
+        let filename = format!(
+            "d{depth:02}_s{stream_id:02X}_r{:03}_id{:08X}_{}.{}",
+            record.index,
+            record.command_id,
+            &digest[..12],
+            spec.suffix
+        );
+        let path = self.output_dir.join(filename);
+        write_atomic(&path, data)?;
+        let artifact = ArtifactReport {
+            kind: spec.kind.to_owned(),
+            path: path
+                .file_name()
+                .ok_or_else(|| Error::Invalid("artifact path has no file name".to_owned()))?
+                .to_string_lossy()
+                .into_owned(),
+            size: data.len(),
+            sha256: digest,
+            depth,
+            stream_id,
+            record_index: Some(record.index),
+            command_id: Some(record.command_id),
+            classification: spec.classification.to_owned(),
+        };
+        self.artifacts.push(artifact.clone());
+        Ok(artifact)
+    }
+}
+
+fn record_report(record: Record) -> RecordReport {
+    RecordReport {
+        index: record.index,
+        command_id: record.command_id,
+        flags: record.flags,
+        image_offset: record.image_offset,
+        image_size: record.image_size,
+        metadata_offset: record.metadata_offset,
+        metadata_size: record.metadata_size,
+        id_copy: record.id_copy,
+        entry_offset: record.entry_offset,
+        init_offset: record.init_offset,
+        direct: record.flags & DIRECT_FLAG != 0,
+        extraction_status: "complete".to_owned(),
+        image: None,
+        metadata: None,
+        nested_stream_id: None,
+    }
+}
+
+fn decoder_report(
+    kind: &str,
+    interpreter_id: Option<u32>,
+    config: &Module9bConfig,
+) -> DecoderReport {
+    DecoderReport {
+        kind: kind.to_owned(),
+        interpreter_id,
+        header_seed: config.header_seed,
+        container_seed: config.container_seed,
+        schedule_offset: config.schedule_offset,
+        aes_key_sha256: sha256(&config.aes_key),
+        skip_aes: config.skip_aes,
+    }
+}
+
+fn record_slice(stream: &[u8], offset: u32, size: u32) -> Result<&[u8]> {
+    let offset = usize::try_from(offset)
+        .map_err(|_| Error::Invalid("record payload offset exceeds usize".to_owned()))?;
+    let size = usize::try_from(size)
+        .map_err(|_| Error::Invalid("record payload size exceeds usize".to_owned()))?;
+    let end = offset
+        .checked_add(size)
+        .ok_or_else(|| Error::Invalid("record payload range overflows usize".to_owned()))?;
+    stream.get(offset..end).ok_or_else(|| {
+        Error::Invalid(format!(
+            "record payload range 0x{offset:x}..0x{end:x} exceeds stream 0x{:x}",
+            stream.len()
+        ))
+    })
+}
+
+fn record_tail(stream: &[u8], offset: u32) -> Result<&[u8]> {
+    let offset = usize::try_from(offset)
+        .map_err(|_| Error::Invalid("record container offset exceeds usize".to_owned()))?;
+    stream.get(offset..).ok_or_else(|| {
+        Error::Invalid(format!(
+            "record container offset 0x{offset:x} exceeds stream 0x{:x}",
+            stream.len()
+        ))
+    })
+}
+
+fn stage1_report(stage1: &Stage1Result, outer_size: usize) -> Stage1Report {
+    Stage1Report {
+        section_index: stage1.section_index,
+        section_type: SHT_LOUSER,
+        section_offset: stage1.section_offset,
+        section_size: stage1.section_size,
+        outer_size,
+        header_offset: stage1.header_offset,
+        header_key: stage1.header.key,
+        payload_offset: stage1.header.payload_offset,
+        payload_size: stage1.header.payload_size,
+        payload_key: stage1.header.payload_key,
+        entry_offset: stage1.header.entry_offset,
+        protect_size: stage1.header.protect_size,
+        stage2_file_offset: stage1.payload_file_offset,
+        stage2_size: stage1.plaintext.len(),
+        stage2_sha256: sha256(&stage1.plaintext),
+        remaining_file_offset: stage1.remaining_file_offset,
+        remaining_size: stage1.remaining_size,
+    }
+}
+
+fn write_json_atomic(path: &Path, value: &impl serde::Serialize) -> Result<()> {
+    let mut bytes = to_vec_pretty(value)?;
+    bytes.push(b'\n');
+    write_atomic(path, &bytes)
+}
+
+fn write_atomic(path: &Path, data: &[u8]) -> Result<()> {
+    let parent = path.parent().unwrap_or_else(|| Path::new("."));
+    create_dir_all(parent)
+        .map_err(|source| Error::io("create output directory", parent, source))?;
+    let mut temporary = NamedTempFile::new_in(parent)
+        .map_err(|source| Error::io("create temporary output", parent, source))?;
+    temporary
+        .write_all(data)
+        .and_then(|()| temporary.as_file().sync_all())
+        .map_err(|source| Error::io("write temporary output", temporary.path(), source))?;
+    temporary
+        .persist(path)
+        .map_err(|error| Error::io("replace output", path, error.error))?;
+    Ok(())
+}
+
+fn absolute(path: &Path) -> Result<PathBuf> {
+    if path.is_absolute() {
+        Ok(path.to_path_buf())
+    } else {
+        std::env::current_dir()
+            .map(|current| current.join(path))
+            .map_err(|source| Error::io("query current directory", path, source))
+    }
+}
+
+fn sha256(data: &[u8]) -> String {
+    let mut digest = Sha256::new();
+    digest.update(data);
+    hex_digest(&digest.finalize())
+}
+/// Lowercase hex of a digest output (sha2 0.11's `Array` no longer formats as
+/// hex directly).
+fn hex_digest(data: &[u8]) -> String {
+    let mut out = String::with_capacity(data.len() * 2);
+    for byte in data {
+        out.push_str(&format!("{byte:02x}"));
+    }
+    out
+}