androscalpel/apk_frauder/src/file_header.rs

use log::warn;
use std::io::{SeekFrom, Write};

use crate::compression::CompressionMethod;
use crate::error::Error;
use crate::extra_fields::{ExtraField, GenericExtraField, Zip64ExtraField};
use crate::{cp437, external_file_attributes, general_purpose_flags, Encoding, Signature};
use androscalpel_serializer::{ReadSeek, Result, Serializable};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct FileHeader {
    // signature: Signature(0x02014b50)
    pub version_made_by: u16,
    pub version_needed_to_extract: u16,
    pub general_purpose_flags: u16,
    pub compression_method: CompressionMethod,
    pub last_mod_file_time: u16,
    pub last_mod_file_data: u16,
    pub crc_32: u32,
    pub compressed_size: u32,
    pub uncompressed_size: u32,
    // file_name_length: u16,
    // extra_field_length: u16,
    // file_comment_length: u16,
    pub disk_number_start: u16,
    pub internal_file_attributes: u16,
    pub external_file_attributes: u32,
    pub offset_local_header: u32,
    pub file_name: Vec<u8>,
    pub extra_field: Vec<ExtraField>,
    /// Remaining bytes in the extra_fields that could not be parsed as ExtraField
    pub malformed_extra_field: Vec<u8>,
    pub file_comment: Vec<u8>,
}

impl Serializable for FileHeader {
    fn serialize(&self, output: &mut dyn Write) -> Result<()> {
        Self::SIGNATURE.serialize(output)?;
        self.version_made_by.serialize(output)?;
        self.version_needed_to_extract.serialize(output)?;
        self.general_purpose_flags.serialize(output)?;
        self.compression_method.serialize(output)?;
        self.last_mod_file_time.serialize(output)?;
        self.last_mod_file_data.serialize(output)?;
        self.crc_32.serialize(output)?;
        self.compressed_size.serialize(output)?;
        self.uncompressed_size.serialize(output)?;
        (self.file_name.len() as u16).serialize(output)?;
        (self
            .extra_field
            .iter()
            .map(|f| f.size() as u16)
            .sum::<u16>()
            + self.malformed_extra_field.len() as u16)
            .serialize(output)?;
        (self.file_comment.len() as u16).serialize(output)?;
        self.disk_number_start.serialize(output)?;
        self.internal_file_attributes.serialize(output)?;
        self.external_file_attributes.serialize(output)?;
        self.offset_local_header.serialize(output)?;
        for c in &self.file_name {
            c.serialize(output)?;
        }
        for c in &self.extra_field {
            c.serialize(output)?;
        }
        for c in &self.malformed_extra_field {
            c.serialize(output)?;
        }
        for c in &self.file_comment {
            c.serialize(output)?;
        }
        Ok(())
    }

    fn deserialize(input: &mut dyn ReadSeek) -> Result<Self> {
        let signature = Signature::deserialize(input)?;
        assert_eq!(signature, Self::SIGNATURE); // TODO
        let version_made_by = u16::deserialize(input)?;
        let version_needed_to_extract = u16::deserialize(input)?;
        let general_purpose_flags = u16::deserialize(input)?;
        let compression_method = CompressionMethod::deserialize(input)?;
        let last_mod_file_time = u16::deserialize(input)?;
        let last_mod_file_data = u16::deserialize(input)?;
        let crc_32 = u32::deserialize(input)?;
        let compressed_size = u32::deserialize(input)?;
        let uncompressed_size = u32::deserialize(input)?;
        let file_name_length = u16::deserialize(input)?;
        let extra_field_length = u16::deserialize(input)?;
        let file_comment_length = u16::deserialize(input)?;
        let disk_number_start = u16::deserialize(input)?;
        let internal_file_attributes = u16::deserialize(input)?;
        let external_file_attributes = u32::deserialize(input)?;
        let offset_local_header = u32::deserialize(input)?;
        let mut file_name = vec![];
        for _ in 0..file_name_length {
            file_name.push(u8::deserialize(input)?);
        }
        let mut header = Self {
            version_made_by,
            version_needed_to_extract,
            general_purpose_flags,
            compression_method,
            last_mod_file_time,
            last_mod_file_data,
            crc_32,
            compressed_size,
            uncompressed_size,
            disk_number_start,
            internal_file_attributes,
            external_file_attributes,
            offset_local_header,
            file_name,
            extra_field: vec![],
            malformed_extra_field: vec![],
            file_comment: vec![],
        };
        //let end_of_extra_field = input.stream_position().unwrap() + extra_field_length as u64;
        let extra_field_off = input.stream_position().unwrap();
        let mut extra_size_read = 0;
        while extra_size_read < extra_field_length as usize {
            let field_off = input.stream_position().unwrap();
            let field = ExtraField::deserialize(input);

            if let Err(err) = field {
                warn!(
                    "Failed to parsed extra field in {}: {err:?}",
                    header.get_name()
                );
                input.seek(SeekFrom::Start(field_off)).unwrap();
                break;
            } else {
                let field = field.unwrap();
                extra_size_read += field.size();
                header.extra_field.push(field);
            }
        }
        if extra_size_read > extra_field_length as usize {
            let last_extra = header.extra_field.pop().unwrap();
            let size = last_extra.size();
            warn!(
                "Failed to parse last extra field in {}, last field ({} bytes long) is {} bytes too long",
                header.get_name(),
                size,
                extra_size_read - extra_field_length as usize
            );
            //warn!("Forgetting last extra field: {:?}", last_extra);
            input.seek(SeekFrom::Current(-(size as i64))).unwrap();
        }
        let mut extra_size_read = input.stream_position().unwrap() - extra_field_off;
        while extra_size_read < extra_field_length as u64 {
            header.malformed_extra_field.push(u8::deserialize(input)?);
            extra_size_read += 1;
        }
        //input.seek(SeekFrom::Start(end_of_extra_field)).unwrap();
        for _ in 0..file_comment_length {
            header.file_comment.push(u8::deserialize(input)?);
        }

        for field in &mut header.extra_field {
            if let ExtraField::Generic(GenericExtraField {
                id: Zip64ExtraField::ID,
                data,
            }) = field
            {
                let original_size = uncompressed_size == u32::MAX;
                let compressed_size = compressed_size == u32::MAX;
                let offset_header = offset_local_header == u32::MAX;
                let disk_number = disk_number_start == u16::MAX;
                let zip64_filed = Zip64ExtraField::from_generic(
                    &GenericExtraField {
                        id: Zip64ExtraField::ID,
                        data: data.clone(),
                    },
                    original_size,
                    compressed_size,
                    offset_header,
                    disk_number,
                )
                .unwrap();
                *field = ExtraField::Zip64(zip64_filed);
            }
        }
        Ok(header)
    }

    fn size(&self) -> usize {
        Self::MIN_SIZE
            + self.file_name.len()
            + self.extra_field.iter().map(|f| f.size()).sum::<usize>()
            + self.malformed_extra_field.len()
            + self.file_comment.len()
    }
}

impl FileHeader {
    const SIGNATURE: Signature = Signature(0x02014b50);
    const MIN_SIZE: usize = 4 + 6 * 2 + 4 * 3 + 5 * 2 + 4 * 2;

    pub fn new_default(name: &str) -> Self {
        let mut header = FileHeader {
            version_made_by: 768,
            version_needed_to_extract: 0,
            general_purpose_flags: 0,
            compression_method: CompressionMethod::Deflated,
            last_mod_file_time: 0,
            last_mod_file_data: 0,
            crc_32: 0,
            compressed_size: 0,
            uncompressed_size: 0,
            disk_number_start: 0,
            internal_file_attributes: 0,
            // TODO: why 0b10000000 ?
            external_file_attributes: external_file_attributes::REGULAR_FILE
                | external_file_attributes::PERM_UR
                | external_file_attributes::PERM_UW
                | external_file_attributes::PERM_GR
                | external_file_attributes::PERM_OR, // TODO
            offset_local_header: 0,
            file_name: vec![],
            extra_field: vec![],
            malformed_extra_field: vec![],
            file_comment: vec![],
        };
        header.set_name(name);
        header
    }

    pub fn get_name_encoding(&self) -> Encoding {
        if self.general_purpose_flags & general_purpose_flags::MASK_UTF8_FILENAME != 0 {
            Encoding::UTF8
        } else {
            Encoding::CP437
        }
    }
    pub fn get_name(&self) -> String {
        match self.get_name_encoding() {
            Encoding::UTF8 => std::str::from_utf8(&self.file_name).unwrap().into(),
            Encoding::CP437 => cp437::cp437_to_string(&self.file_name),
        }
    }
    pub fn external_file_attributes_set_flag(&mut self, flag: u32) {
        self.external_file_attributes |= flag;
    }
    pub fn external_file_attributes_unset_flag(&mut self, flag: u32) {
        self.external_file_attributes &= !flag;
    }
    pub fn set_file_type(&mut self, file_type: u32) {
        self.external_file_attributes &= !external_file_attributes::MASK_FILE_TYPE;
        self.external_file_attributes |= file_type;
    }

    pub fn set_name(&mut self, name: &str) {
        let file_name = match cp437::string_to_cp437(name) {
            Ok(name) => {
                self.general_purpose_flags &= !general_purpose_flags::MASK_UTF8_FILENAME;
                name
            }
            Err(Error::NotCp437) => {
                self.general_purpose_flags |= general_purpose_flags::MASK_UTF8_FILENAME;
                name.as_bytes().into()
            }
        };
        self.file_name = file_name;
    }

    pub fn get_uncompressed_size(&self) -> u64 {
        if self.uncompressed_size != u32::MAX {
            self.uncompressed_size as u64
        } else if let Some(ExtraField::Zip64(Zip64ExtraField {
            original_size: Some(original_size),
            ..
        })) = self
            .extra_field
            .iter()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *original_size
        } else {
            self.uncompressed_size as u64
        }
    }

    pub fn set_uncompressed_size(&mut self, uncompressed_size: u64) {
        if let Some(ExtraField::Zip64(Zip64ExtraField {
            original_size: Some(original_size),
            ..
        })) = self
            .extra_field
            .iter_mut()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *original_size = uncompressed_size;
            self.uncompressed_size = u32::MAX;
        } else if uncompressed_size > u32::MAX as u64 {
            self.extra_field.push(ExtraField::Zip64(Zip64ExtraField {
                original_size: Some(uncompressed_size),
                compressed_size: Some(self.compressed_size as u64),
                disk_number: None,
                offset_header: None,
            }));
            self.uncompressed_size = u32::MAX;
            self.compressed_size = u32::MAX;
        } else {
            self.uncompressed_size = uncompressed_size as u32;
        }
    }

    pub fn get_compressed_size(&self) -> u64 {
        if self.compressed_size != u32::MAX {
            self.compressed_size as u64
        } else if let Some(ExtraField::Zip64(Zip64ExtraField {
            compressed_size: Some(compressed_size),
            ..
        })) = self
            .extra_field
            .iter()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *compressed_size
        } else {
            self.compressed_size as u64
        }
    }

    pub fn set_compressed_size(&mut self, compressed_size: u64) {
        if let Some(ExtraField::Zip64(Zip64ExtraField {
            compressed_size: Some(compressed_size_),
            ..
        })) = self
            .extra_field
            .iter_mut()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *compressed_size_ = compressed_size;
            self.compressed_size = u32::MAX;
        } else if compressed_size > u32::MAX as u64 {
            self.extra_field.push(ExtraField::Zip64(Zip64ExtraField {
                original_size: Some(self.uncompressed_size as u64),
                compressed_size: Some(compressed_size),
                disk_number: None,
                offset_header: None,
            }));
            self.uncompressed_size = u32::MAX;
            self.compressed_size = u32::MAX;
        } else {
            self.compressed_size = compressed_size as u32;
        }
    }

    pub fn get_offset_local_header(&self) -> u64 {
        if self.offset_local_header != u32::MAX {
            self.offset_local_header as u64
        } else if let Some(ExtraField::Zip64(Zip64ExtraField {
            offset_header: Some(offset_header),
            ..
        })) = self
            .extra_field
            .iter()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *offset_header
        } else {
            self.offset_local_header as u64
        }
    }

    pub fn set_offset_local_header(&mut self, offset: u64) {
        if let Some(ExtraField::Zip64(Zip64ExtraField {
            offset_header: Some(offset_header),
            ..
        })) = self
            .extra_field
            .iter_mut()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *offset_header = offset;
            self.offset_local_header = u32::MAX;
        } else if offset > u32::MAX as u64 {
            self.extra_field.push(ExtraField::Zip64(Zip64ExtraField {
                original_size: Some(self.uncompressed_size as u64),
                compressed_size: Some(self.compressed_size as u64),
                disk_number: None,
                offset_header: Some(offset),
            }));
            self.uncompressed_size = u32::MAX;
            self.compressed_size = u32::MAX;
            self.offset_local_header = u32::MAX;
        } else {
            self.offset_local_header = offset as u32;
        }
    }

    pub fn get_disk_number_start(&self) -> u32 {
        if self.disk_number_start != u16::MAX {
            self.disk_number_start as u32
        } else if let Some(ExtraField::Zip64(Zip64ExtraField {
            disk_number: Some(disk_number),
            ..
        })) = self
            .extra_field
            .iter()
            .find(|f| matches!(f, ExtraField::Zip64(_)))
        {
            *disk_number
        } else {
            self.disk_number_start as u32
        }
    }
}