import logging

from malduck import asciiz
from malduck.extractor import Extractor
from malduck.pe import PE
from malduck.procmem import ProcessMemory
from malduck.yara import YaraRuleMatch

from ...utils import Config, get_rule_metadata

log = logging.getLogger(__name__)

__author__ = "4rchib4ld"
__version__ = "1.0.0"


class IcedIDPhotoLoader(Extractor):

    """
    IcedID PhotoLoader C2 Domain Configuration Extractor
    """

    family: str = "icedid_photoloader"
    yara_rules: tuple = ("icedid_photoloader",)

    @staticmethod
    def extractPayload(pe: PE) -> bytes:
        """
        Extracting the payload from the .data section
        """
        for section in pe.sections:
            if ".data" in str(section.Name):
                data = section.get_data()
                payload = asciiz(data[4:])
                return payload

        return b""

    @Extractor.rule
    def icedid_photoloader(self, p: ProcessMemory, match: YaraRuleMatch) -> Config | bool:
        conf: Config = {}
        info: Config = get_rule_metadata(match)

        obfuscationCode = match.elements["obfuscationCode"][0][2]
        xorCountValue = obfuscationCode[3]  # Getting this values dynamically because... you never know
        countValue = obfuscationCode[-1]
        pe_rep = PE(data=p)
        payload = self.extractPayload(pe_rep)
        decrypted = bytearray()
        for i in range(countValue):
            try:
                decrypted.append(payload[i + xorCountValue] ^ payload[i])
            except IndexError:
                pass

        c2 = asciiz(decrypted)
        if len(c2) > 0:
            conf[self.family] = [c2.decode("utf-8")]

        return conf | info