#!/usr/bin/python3
# -*- coding: utf-8 -*-
# (c) B.Kerler 2018-2024 GPLv3 License
import logging
import time
from struct import pack, unpack

from mtkclient.Library.Hardware.hwcrypto import HwCrypto, CryptoSetup
from mtkclient.Library.utils import LogBase, logsetup
from mtkclient.config.payloads import PathConfig


class Exploitation(metaclass=LogBase):
    def __init__(self, mtk, loglevel=logging.INFO):
        self.__logger, self.info, self.debug, self.warning, self.error = logsetup(self, self.__logger, 
                                                                                  loglevel, mtk.config.gui)
        self.lasterror = ""
        self.mtk = mtk
        self.chipconfig = self.mtk.config.chipconfig
        self.config = self.mtk.config
        self.usbwrite = self.mtk.port.usbwrite
        self.usbread = self.mtk.port.usbread
        self.read32 = self.mtk.preloader.read32
        self.write32 = self.mtk.preloader.write32
        self.hwcode = mtk.config.hwcode

        # crypto types
        setup = CryptoSetup()
        setup.hwcode = self.mtk.config.hwcode
        setup.dxcc_base = self.mtk.config.chipconfig.dxcc_base
        setup.read32 = self.mtk.preloader.read32
        setup.write32 = self.mtk.preloader.write32
        setup.writemem = self.mtk.preloader.writemem
        setup.da_payload_addr = self.mtk.config.chipconfig.da_payload_addr
        setup.gcpu_base = self.mtk.config.chipconfig.gcpu_base
        setup.blacklist = self.mtk.config.chipconfig.blacklist
        setup.sej_base = self.mtk.config.chipconfig.sej_base
        setup.cqdma_base = self.mtk.config.chipconfig.cqdma_base
        setup.ap_dma_mem = self.mtk.config.chipconfig.ap_dma_mem
        setup.meid_addr = self.mtk.config.chipconfig.meid_addr
        setup.prov_addr = self.mtk.config.chipconfig.prov_addr
        self.hwcrypto = HwCrypto(setup, loglevel, self.mtk.config.gui)

        self.chipconfig = self.mtk.config.chipconfig
        self.var1 = self.chipconfig.var1
        self.pathconfig = PathConfig()

    def fix_payload(self, payload, da=True):
        payload = bytearray(payload)
        wd = unpack("<I", payload[-4:])[0]
        ua = unpack("<I", payload[-8:-4])[0]
        if wd == 0x10007000:
            payload[-4:] = pack("<I", self.mtk.config.chipconfig.watchdog)
        if ua == 0x11002000:
            payload[-8:-4] = pack("<I", self.mtk.config.chipconfig.uart)
        while len(payload) % 4 != 0:
            payload.append(0)
        if da:
            payload.extend(b"\x00" * 0x100)  # signature len
        return payload

    def da_read(self, address, length, check_result=True):
        return self.da_read_write(address, length, None, check_result)

    def da_write(self, address, length, data, check_result=True):
        return self.da_read_write(address, length, data, check_result)

    def exploit(self, payload, payloadaddr):
        pass

    def da_read_write(self, address, length, data=None, check_result=True):
        pass

    def da_payload(self, payload, addr, forcekamakiri=True, exploittype=1):
        hassecurity = self.mtk.config.target_config["sla"] or self.mtk.config.target_config["daa"] or \
                      self.mtk.config.target_config["sbc"]
        is_v6 = self.mtk.config.chipconfig.damode == 6
        if is_v6:
            forcekamakiri = False
        if hassecurity or forcekamakiri:
            try:
                payload = self.fix_payload(bytearray(payload), False)
            except Exception:
                pass

            try:
                if self.exploit(payload, addr):
                    self.info("Done sending payload...")
                    time.sleep(0.2)
                    return True
            except Exception:
                self.error("Error on sending payload.")
                return False
        else:
            self.info("Sending payload via insecure da.")
            payloadaddr = 0x200000
            if self.chipconfig.dacode in self.mtk.daloader.daconfig.dasetup:
                entry = self.mtk.daloader.daconfig.dasetup[self.chipconfig.dacode][0]
                payloadaddr = entry.region[1].m_start_addr
            payload = self.fix_payload(payload, True)
            if self.mtk.preloader.send_da(payloadaddr, len(payload) - 0x100, 0x100, payload):
                if self.mtk.preloader.jump_da(payloadaddr):
                    self.info("Done sending payload...")
                    return True
            self.error("Error on sending payload.")
            return False

    def bruteforce(self, args, startaddr=0x9900):
        pass

    def newbrute(self, dump_ptr, dump=False):
        pass

    def dump_preloader(self, filename=None):
        pass

    def dump_brom(self, filename, dump_ptr=None, length=0x20000):
        pass

    def payload(self, payload, daaddr):
        pass

    def runpayload(self, payload, ack=0xA1A2A3A4, addr=None, dontack=False):
        self.info("Running payload")
        if addr is None:
            if self.mtk.config.target_config["sla"] or self.mtk.config.target_config["daa"]:
                addr = self.chipconfig.brom_payload_addr
            else:
                addr = self.chipconfig.da_payload_addr
            if self.da_payload(payload, addr, False):
                if dontack:
                    return True
                result = self.usbread(4)
                if result == pack(">I", ack):
                    self.info("Successfully sent payload.")
                    return True
            self.error("Error on sending payload.")
        return False

    def close(self):
        pass

    def crash(self, mode=0):
        self.info("Crashing da...")
        try:
            if mode == 0:
                self.mtk.preloader.send_da(0, 0x100, 0x100, b'\x00' * 0x100)
            elif mode == 1:
                self.mtk.preloader.read32(0, 0x100)
            elif mode == 2:
                payload = b'\x00\x01\x9F\xE5\x10\xFF\x2F\xE1' + b'\x00' * 0x110
                self.mtk.preloader.send_da(0x0, len(payload), 0x0, payload)
                self.mtk.preloader.jump_da(0x0)
        except Exception:
            pass
