Kasra Rafiezadeh Shahi

import logging
import zipfile
from abc import ABCMeta, abstractmethod

from redb.extractors.extractor import Extractor

logger = logging.getLogger(__name__)


@abstractmethod
class APKExtractor(Extractor, metaclass=ABCMeta):

    def __init__(
        self,
        filepath,
        log,
        exporters=None,
        index_prefix=None,
        known_benign=False,
        known_malicious=False,
        apk=None,
    ):
        # Read binary data before calling super().__init__ so that the
        # binary property override is available during hash computation.
        with open(filepath, "rb") as f:
            self._binary_data = f.read()

        self.apk = apk if apk else self._generate_apk_object(filepath)

        super().__init__(
            filepath,
            log,
            exporters,
            index_prefix,
            known_benign=known_benign,
            known_malicious=known_malicious,
        )

    @property
    def binary(self):
        """Override to use already-read binary data."""
        return self._binary_data

    def _generate_apk_object(self, filepath):
        """Parse APK using androguard."""
        apk = None
        try:
            # Suppress verbose androguard loguru DEBUG/INFO output
            from loguru import logger as loguru_logger
            loguru_logger.disable("androguard")
            from androguard.core.apk import APK
            apk = APK(filepath)
            if not apk.is_valid_APK():
                logger.warning(
                    f"APK validation warning for {filepath}"
                )
                # Still return the object — partial parsing may still work
        except Exception as e:
            logger.error(
                f"Format error parsing APK {filepath}: {e}"
            )
        return apk

    def _is_valid_apk(self):
        """Check if the APK object was parsed successfully."""
        return self.apk is not None

    def _get_zip_file(self):
        """Get a zipfile.ZipFile handle for direct archive inspection."""
        try:
            return zipfile.ZipFile(self.filepath, 'r')
        except (zipfile.BadZipFile, Exception) as e:
            self.log.error(f"Failed to open APK as ZIP: {e}")
            return None

    def _list_files(self):
        """List all files in the APK archive."""
        if not self._is_valid_apk():
            return []
        try:
            return self.apk.get_files()
        except Exception as e:
            self.log.error(f"Error listing APK files: {e}")
            return []