# -*- coding:utf-8 -*-
# Author: Kei Choi(hanul93@gmail.com)

"""
eMalware Engine Plugin

This plugin handles malware detection for PE and ELF executables.
"""

import logging
import os
import re

from kicomav.plugins import cryptolib
from kicomav.plugins import kavutil
from kicomav.plugins import kernel
from kicomav.plugins import pe
from kicomav.kavcore import k2security
from kicomav.kavcore.plugin_base import MalwareDetectorBase

# Module logger
logger = logging.getLogger(__name__)


# -------------------------------------------------------------------------
# KavMain Class
# -------------------------------------------------------------------------
class KavMain(MalwareDetectorBase):
    """eMalware detector plugin.

    This plugin provides functionality for:
    - Detecting malware in PE files (sections, attach area, PDB info)
    - Detecting malware in ELF files (sections, program headers)
    - Detecting Mirai variants in ELF files
    """

    def __init__(self):
        """Initialize the eMalware Engine plugin."""
        super().__init__(
            author="Kei Choi",
            version="1.1",
            title="eMalware Engine",
            kmd_name="emalware",
        )
        self.p_allaple = None
        self.p_linux_mirai = None
        self.mirai_a_strings = []
        self.aho_mirai_a = None

    def _load_virus_database(self) -> int:
        """Load virus patterns.

        Returns:
            0 for success
        """
        # Worm.Win32.Allaple.gen pattern
        pat = rb"\xC7\x44\x24.+[\x40\x41]\x00.+\x8B.\x24.+\x33..\x04\x01.\x24.+\x8B.\x24"
        self.p_allaple = re.compile(pat)

        pat = (
            rb"POST /cdn-cgi/\x00\x00 HTTP/1.1\r\nUser-Agent: \x00\r\nHost:"
            + rb"[\d\D]+?GET\x00+/\x00+Cookie:[\d\D]+?http[\d\D]+?url=[\d\D]+?POST"
        )
        self.p_linux_mirai = re.compile(pat)

        # Mirai variant virus pattern
        self.mirai_a_strings = [
            b"POST /cdn-cgi/",
            b"HTTP/1.1\r\nUser-Agent: ",
            b"Host:",
            b"GET",
            b"Cookie:",
            b"http",
            b"url=",
            b"proc/net/tcp",
        ]

        self.aho_mirai_a = kavutil.AhoCorasick()
        self.aho_mirai_a.make_tree(self.mirai_a_strings)

        return 0

    def getinfo(self):
        """Get plugin information.

        Returns:
            Dictionary containing plugin metadata
        """
        info = super().getinfo()
        s_num = 5
        if kavutil.handle_pattern_md5:
            s_num = kavutil.handle_pattern_md5.get_sig_num("emalware") + 5
        info["sig_num"] = s_num
        return info

    def listvirus(self):
        """Get list of detectable viruses.

        Returns:
            List of virus names
        """
        vlists = kavutil.handle_pattern_md5.get_sig_vlist("emalware")
        if vlists is None:
            vlists = []
        vlists.append("Backdoor.Linux.Mirai.a.gen")
        vlists.append(kavutil.normal_vname("<n>AdWare.Win32.Sokuxuan.gen"))
        vlists = sorted(set(vlists))

        return vlists

    def scan(self, filehandle, filename, fileformat, filename_ex):
        """Scan for malware.

        Args:
            filehandle: File data (memory mapped)
            filename: Path to file
            fileformat: Format info from format() method
            filename_ex: Extended filename info

        Returns:
            Tuple of (found, malware_name, malware_id, result)
        """
        try:
            mm = filehandle

            # Check if the file format is PE and not None
            if fileformat and "ff_pe" in fileformat:
                ret = self.__scan_pe(mm, fileformat)
                if ret[0]:
                    return ret

            elif fileformat and "ff_elf" in fileformat:
                ret = self.__scan_elf(mm, fileformat)
                if ret[0]:
                    return ret

            # For NSIS-like installers, the attached area contains the original image.
            if "ff_attach" in fileformat:
                foff = fileformat["ff_attach"]["Attached_Pos"]
                buf = mm[foff:]
                fsize = len(buf)
                if fsize and kavutil.handle_pattern_md5.match_size("emalware", fsize):
                    fmd5 = cryptolib.md5(buf)
                    if vname := kavutil.handle_pattern_md5.scan("emalware", fsize, fmd5):
                        vname = kavutil.normal_vname(vname)
                        return True, vname, kernel.DISINFECT_DELETE, kernel.INFECTED

        except (IOError, OSError) as e:
            logger.debug("Scan IO error for %s: %s", filename, e)
        except Exception as e:
            logger.warning("Unexpected error scanning %s: %s", filename, e)

        return False, "", kernel.DISINFECT_NONE, kernel.NOT_FOUND

    def __scan_pe(self, mm, fileformat):
        """Scan PE files for malware."""
        ff = fileformat["ff_pe"]

        # case 1 : Scan the entire section by hash
        for section in ff["pe"]["Sections"]:
            fsize = section["SizeRawData"]
            if fsize and kavutil.handle_pattern_md5.match_size("emalware", fsize):
                foff = section["PointerRawData"]
                fmd5 = cryptolib.md5(mm[foff : foff + fsize])
                if vname := kavutil.handle_pattern_md5.scan("emalware", fsize, fmd5):
                    vname = kavutil.normal_vname(vname)
                    return True, vname, kernel.DISINFECT_DELETE, kernel.INFECTED

        # case 2. Check if the executable file exists in the last section
        if len(ff["pe"]["Sections"]):
            sec = ff["pe"]["Sections"][-1]
            off = sec["PointerRawData"]
            size = sec["SizeRawData"]

            exe_offs = [m.start() for m in re.finditer(b"MZ", mm[off : off + size])]

            for exe_pos in exe_offs:
                fsize = 0x1D5
                if fsize and kavutil.handle_pattern_md5.match_size("emalware", fsize):
                    fmd5 = cryptolib.md5(mm[off + exe_pos : off + exe_pos + fsize])
                    if vname := kavutil.handle_pattern_md5.scan("emalware", fsize, fmd5):
                        idx = len(ff["pe"]["Sections"]) - 1
                        vname = kavutil.normal_vname(vname)
                        return True, vname, (kernel.DISINFECT_MALWARE | idx), kernel.INFECTED

        # case 3. Check for malware using pdb
        if "PDB_Name" in ff["pe"]:
            pdb_sigs = {
                b":\\pz_git\\bin\\": "<n>AdWare.Win32.Sokuxuan.gen",
                b":\\CODE\\vitruvian\\": "<n>AdWare.Win32.Vitruvian.gen",
                b"\\bin\\Release\\WebSparkle.": "<n>AdWare.MSIL.BrowseFox.gen",
                b":\\TeamCity\\BuildAgent1\\work\\": "<n>WebToolbar.Win32.Agent.avi",
            }

            for pat in pdb_sigs:
                if pat in ff["pe"]["PDB_Name"]:
                    vname = kavutil.normal_vname(pdb_sigs[pat])
                    return True, vname, kernel.DISINFECT_DELETE, kernel.INFECTED

        # case 4. Check for Worm.Win32.Allaple.gen
        ep_off = ff["pe"]["EntryPointRaw"]
        data = mm[ep_off : ep_off + 0x80]
        if self.p_allaple and self.p_allaple.search(data):
            return True, "Worm.Win32.Allaple.gen", kernel.DISINFECT_DELETE, kernel.INFECTED

        return False, "", kernel.DISINFECT_NONE, kernel.NOT_FOUND

    def __scan_elf(self, mm, fileformat):
        """Scan ELF files for malware."""
        ff = fileformat["ff_elf"]
        elf_sections = ff["elf"].get("Sections", [])
        elf_prog_headers = ff["elf"].get("ProgramHeaders", [])

        # Case 5: Check if the executable file exists in the sections
        if len(elf_sections):
            for section in elf_sections:
                if (section["Type"] & 0x1) == 0x1 and (section["Flag"] & 0x4) == 0x4:
                    fsize = section["Size"]
                    if fsize and kavutil.handle_pattern_md5.match_size("emalware", fsize):
                        foff = section["Offset"]
                        fmd5 = cryptolib.md5(mm[foff : foff + fsize])
                        if vname := kavutil.handle_pattern_md5.scan("emalware", fsize, fmd5):
                            vname = kavutil.normal_vname(vname)
                            return True, vname, kernel.DISINFECT_DELETE, kernel.INFECTED

        # Case 6: Check if the executable file exists in the program headers
        elif len(elf_prog_headers):
            for ph in elf_prog_headers:
                if (ph["Type"] & 0x1) == 0x1 and (ph["Flag"] & 0x1) == 0x1:
                    fsize = ph["Size"]
                    if fsize and kavutil.handle_pattern_md5.match_size("emalware", fsize):
                        foff = ph["Offset"]
                        fmd5 = cryptolib.md5(mm[foff : foff + fsize])
                        if vname := kavutil.handle_pattern_md5.scan("emalware", fsize, fmd5):
                            vname = kavutil.normal_vname(vname)
                            return True, vname, kernel.DISINFECT_DELETE, kernel.INFECTED

        # Case 7. Check for Mirai variant
        for section in elf_sections:
            if section["Name"] == ".rodata":
                vstring = []

                foff = section["Offset"]
                ret = self.aho_mirai_a.search(mm[foff : foff + 0x200])

                vstring.extend(n[1] for n in ret[: len(self.mirai_a_strings)])

                if set(vstring) == set(self.mirai_a_strings):
                    return True, "Backdoor.Linux.Mirai.a.gen", kernel.DISINFECT_DELETE, kernel.INFECTED

        return False, "", kernel.DISINFECT_NONE, kernel.NOT_FOUND

    def disinfect(self, filename, malware_id):
        """Disinfect malware.

        Args:
            filename: Path to infected file
            malware_id: Malware ID to disinfect

        Returns:
            True if successful, False otherwise
        """
        try:
            if malware_id == kernel.DISINFECT_DELETE:
                filename_dir = os.path.dirname(filename) or os.getcwd()
                k2security.safe_remove_file(filename, filename_dir)
                return True

            if malware_id & kernel.DISINFECT_MALWARE == kernel.DISINFECT_MALWARE:
                idx = malware_id & ~kernel.DISINFECT_MALWARE

                with open(filename, "rb") as f:
                    buf = f.read()

                pe_ = pe.PE(buf, False, filename)
                try:
                    pe_format = pe_.parse()
                except MemoryError:
                    pe_format = None

                if pe_format is None:
                    return False

                ff = {"pe": pe_format}

                if len(ff["pe"]["Sections"]) > idx:
                    return self.clean_pe_section(ff, idx, buf, filename)

        except (IOError, OSError, k2security.SecurityError) as e:
            logger.debug("Disinfect error for %s: %s", filename, e)
        except Exception as e:
            logger.warning("Unexpected error disinfecting %s: %s", filename, e)

        return False

    def clean_pe_section(self, ff, idx, buf, filename):
        """Clean up a PE section."""
        section = ff["pe"]["Sections"][idx]
        fsize = section["SizeRawData"]
        foff = section["PointerRawData"]

        data = buf[:foff] + (b"\x00" * fsize) + buf[foff + fsize :]
        with open(filename, "wb") as f:
            f.write(data)
        return True
