from __future__ import annotations
from typing import TYPE_CHECKING

from collections import deque

from trueseeing.core.model.cmd import CommandMixin
from trueseeing.core.ui import ui

if TYPE_CHECKING:
  from typing import Any, Optional
  from trueseeing.api import CommandHelper, Command, CommandMap

class ExploitCommand(CommandMixin):
  def __init__(self, helper: CommandHelper) -> None:
    self._helper = helper

  @staticmethod
  def create(helper: CommandHelper) -> Command:
    return ExploitCommand(helper)

  def get_commands(self) -> CommandMap:
    return {
      'xq':dict(e=self._exploit_discard, n='xq', d='exploit: discard changes'),
      'xx':dict(e=self._exploit_apply, n='xx[!]', d='exploit: apply and rebuild apk'),
      'xx!':dict(e=self._exploit_apply),
      #'xf':dict(e=self._exploit_inject, n='xf', d='exploit; inject frida gadget'),
      'xu':dict(e=self._exploit_disable_pinning, n='xu', d='exploit: disable SSL/TLS pinning'),
      'xd':dict(e=self._exploit_enable_debug, n='xd', d='exploit: make debuggable'),
      'xb':dict(e=self._exploit_enable_backup, n='xb', d='exploit: make backupable'),
      'xt':dict(e=self._exploit_patch_target_api_level, n='xt[!] <api level>', d='exploit: patch target api level'),
      'xt!':dict(e=self._exploit_patch_target_api_level),
      'xp':dict(e=self._exploit_device_list_packages, n='xp', d='device: list installed packages'),
      'xco':dict(e=self._exploit_device_copyout, n='xco[!] package [data.tar]', d='device: copy-out package data'),
      'xco!':dict(e=self._exploit_device_copyout),
      'xci':dict(e=self._exploit_device_copyin, n='xci[!] package [data.tar]', d='device: copy-in package data'),
      'xci!':dict(e=self._exploit_device_copyin),
    }

  async def _exploit_discard(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    _ = args.popleft()

    import time

    at = time.time()

    context = await self._helper.get_context().require_type('apk').analyze(level=2)
    with context.store().query().scoped() as q:
      if not q.patch_exists(None):
        ui.fatal('nothing to discard')
      ui.info('discarding patches to {apk}'.format(apk=apk))
      q.patch_clear()

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_apply(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    cmd = args.popleft()

    import os
    import time
    from tempfile import TemporaryDirectory
    from trueseeing.core.android.asm import APKAssembler
    from trueseeing.core.android.tools import move_apk

    origapk = apk.replace('.apk', '.apk.orig')

    if os.path.exists(origapk) and not cmd.endswith('!'):
      ui.fatal('backup file exists; force (!) to overwrite')

    at = time.time()

    context = await self._helper.get_context().require_type('apk').analyze(level=2)
    with context.store().query().scoped() as q:
      if not q.patch_exists(None):
        ui.fatal('nothing to apply')

      with TemporaryDirectory(dir=context.wd) as td:
        ui.info('applying patches to {apk}'.format(apk=apk))
        root = os.path.join(td, 'f')

        for path,blob in q.file_enum(None, patched=True):
          target = os.path.join(root, *path.split('/'))
          os.makedirs(os.path.dirname(target), exist_ok=True)
          with open(target, 'wb') as f:
            f.write(blob)

        outapk, outsig = await APKAssembler.assemble_from_path(td, root)

        if os.path.exists(apk):
          move_apk(apk, origapk)

        move_apk(outapk, apk)

      q.patch_clear()

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_disable_pinning(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    _ = args.popleft()

    import time
    import random
    from importlib.resources import files

    ui.info('disabling declarative TLS pinning {apk}'.format(apk=apk))

    at = time.time()
    context = await self._helper.get_context().require_type('apk').analyze(level=2)
    with context.store().query().scoped() as q:
      key = 'nsc{:04x}'.format(random.randint(0, 2**16))

      path = 'AndroidManifest.xml'
      blob = q.file_get(path, patched=True)
      assert blob is not None

      manif = self._parsed_manifest(blob)
      for e in manif.xpath('.//application'):
        e.attrib['{http://schemas.android.com/apk/res/android}usesCleartextTraffic'] = "true"
        e.attrib['{http://schemas.android.com/apk/res/android}networkSecurityConfig'] = f'@xml/{key}'
      q.patch_put(path, self._manifest_as_xml(manif), z=True)

      # XXX
      path = f'resources/package_1/res/xml/{key}.xml'
      q.patch_put(path, (files('trueseeing')/'libs'/'android'/'nsc.xml').read_bytes(), z=True)

      # XXX
      import lxml.etree as ET
      path = 'resources/package_1/res/values/public.xml'
      root = ET.fromstring(q.file_get(path, patched=True), parser=ET.XMLParser(recover=True))
      if root.xpath('./public[@type="xml"]'):
        maxid = max(int(e.attrib["id"], 16) for e in root.xpath('./public[@type="xml"]'))
        n = ET.SubElement(root, 'public')
        n.attrib['id'] = f'0x{maxid+1:08x}'
        n.attrib['type'] = 'xml'
        n.attrib['name'] = key
      else:
        maxid = (max(int(e.attrib["id"], 16) for e in root.xpath('./public')) & 0xffff0000)
        n = ET.SubElement(root, 'public')
        n.attrib['id'] = f'0x{maxid+0x10000:08x}'
        n.attrib['type'] = 'xml'
        n.attrib['name'] = key
      q.patch_put(path, ET.tostring(root),z=True)

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_enable_debug(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    _ = args.popleft()

    import time

    ui.info('enabling debug {apk}'.format(apk=apk))

    at = time.time()
    context = await self._helper.get_context().require_type('apk').analyze(level=2)
    with context.store().query().scoped() as q:
      path = 'AndroidManifest.xml'
      blob = q.file_get(path, patched=True)
      assert blob is not None
      manif = self._parsed_manifest(blob)
      for e in manif.xpath('.//application'):
        e.attrib['{http://schemas.android.com/apk/res/android}debuggable'] = "true"
      q.patch_put(path, self._manifest_as_xml(manif), z=True)

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_enable_backup(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    _ = args.popleft()

    import time

    ui.info('enabling full backup {apk}'.format(apk=apk))

    at = time.time()
    context = await self._helper.get_context().require_type('apk').analyze(level=1)
    with context.store().query().scoped() as q:
      path = 'AndroidManifest.xml'
      blob = q.file_get(path, patched=True)
      assert blob is not None
      manif = self._parsed_manifest(blob)
      for e in manif.xpath('.//application'):
        e.attrib['{http://schemas.android.com/apk/res/android}allowBackup'] = "true"
        if '{http://schemas.android.com/apk/res/android}fullBackupContent' in e.attrib:
          del e.attrib['{http://schemas.android.com/apk/res/android}fullBackupContent']
      q.patch_put(path, self._manifest_as_xml(manif), z=True)

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_patch_target_api_level(self, args: deque[str]) -> None:
    apk = self._helper.require_target()

    cmd = args.popleft()

    try:
      level = int(args.popleft())
    except (IndexError, ValueError):
      ui.fatal('need API level')

    import time

    ui.info('retargetting API level {level} {apk}'.format(level=level, apk=apk))

    at = time.time()

    context = await self._helper.get_context().require_type('apk').analyze(level=2)
    with context.store().query().scoped() as q:
      path = 'AndroidManifest.xml'
      blob = q.file_get(path, patched=True)
      assert blob is not None
      manif = self._parsed_manifest(blob)
      for e in manif.xpath('.//uses-sdk'):
        e.attrib['{http://schemas.android.com/apk/res/android}targetSdkVersion'] = str(level)
        minLevel = int(e.attrib.get('{http://schemas.android.com/apk/res/android}minSdkVersion', '1'))
        if level < minLevel:
          if not cmd.endswith('!'):
            ui.fatal('cannot target API level below requirement ({minlv}); force (!) to downgrade altogether'.format(minlv=minLevel))
          else:
            ui.warn('downgrading the requirement')
            e.attrib['{http://schemas.android.com/apk/res/android}minSdkVersion'] = str(level)
      q.patch_put(path, self._manifest_as_xml(manif), z=True)

    ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))

  async def _exploit_device_list_packages(self, args: deque[str]) -> None:
    _ = args.popleft()

    ui.info('listing packages')

    import time
    import re
    from trueseeing.core.android.device import AndroidDevice

    at = time.time()
    nr = 0
    for m in re.finditer(r'^package:(.*)', await AndroidDevice().invoke_adb('shell pm list package'), re.MULTILINE):
      p = m.group(1)
      ui.info(p)
      nr += 1
    ui.success('done, {nr} packages found ({t:.02f} sec.)'.format(nr=nr, t=(time.time() - at)))

  async def _exploit_device_copyout(self, args: deque[str]) -> None:
    success: bool = False

    cmd = args.popleft()
    if not args:
      ui.fatal('need package name')

    target = args.popleft()

    import os
    if not args:
      outfn = f'{target}.tar'
    else:
      outfn = args.popleft()

    outfn0 = outfn.replace('.tar', '') + '-int.tar'
    outfn1 = outfn.replace('.tar', '') + '-ext.tar'

    if os.path.exists(outfn) and not cmd.endswith('!'):
      ui.fatal('outfile exists; force (!) to overwrite')

    ui.info(f'copying out: {target} -> {outfn}')

    import time
    from subprocess import CalledProcessError
    from trueseeing.core.tools import invoke_passthru
    from trueseeing.core.android.device import AndroidDevice
    from trueseeing.core.android.tools import toolchains

    at = time.time()
    dev = AndroidDevice()

    if not await dev.is_fullbackup_available():
      ui.warn('full backup feature is not available')
    else:
      ui.info('initiating a backup on device; give "1" as the password if asked')
      await dev.invoke_adb_passthru(f'backup -f {outfn}.ab {target}')
      try:
        try:
          with toolchains() as tc:
            await invoke_passthru('java -jar {abe} unpack {outfn}.ab {outfn} 1'.format(
              abe=tc['abe'],
              outfn=outfn,
            ))
        except CalledProcessError:
          ui.warn('unpack failed (did you give the correct password?); trying the next method')
        else:
          ui.success('unpack success')
          if os.stat(outfn).st_size > 1024:
            ui.success(f'copied out: {outfn}')
            success = True
          else:
            ui.warn('got an empty backup; trying the next method')
            try:
              os.remove(outfn)
            except FileNotFoundError:
              pass
      finally:
        try:
          os.remove(f'{outfn}.ab')
        except FileNotFoundError:
          pass

    if not success:
      if not await dev.is_package_debuggable(target):
        ui.warn('target is not debuggable')
      else:
        ui.info('target seems debuggable; trying extraction with debug interface')

        tfn0 = self._generate_tempfilename_for_device()
        ui.info('copying internal storage')
        await dev.invoke_adb_passthru(f'shell "run-as {target} tar -cv . > {tfn0}"')
        await dev.invoke_adb_passthru(f'pull {tfn0} {outfn0}')
        await dev.invoke_adb_passthru(f'shell rm -f {tfn0}')
        ui.success(f'copied out: {outfn0}')

        ui.info('copying external storage')
        tfn1 = self._generate_tempfilename_for_device()
        try:
          await dev.invoke_adb_passthru(f'shell "cd /storage/emulated/0/Android/ && tar -cv data/{target} obb/{target} > {tfn1}"')
        except CalledProcessError:
          ui.warn('detected errors during extraction from external storage (may indicate partial extraction)')
        await dev.invoke_adb_passthru(f'pull {tfn1} {outfn1}')
        await dev.invoke_adb_passthru(f'shell rm -f {tfn1}')
        ui.success(f'copied out: {outfn1}')

        success = True

    if success:
      ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))
    else:
      ui.failure('copyout failed')

  async def _exploit_device_copyin(self, args: deque[str]) -> None:
    success: bool = False

    _ = args.popleft()
    if not args:
      ui.fatal('need package name')

    target = args.popleft()

    import os
    if not args:
      fn = f'{target}.tar'
    else:
      fn = args.popleft()

    fn0 = fn.replace('.tar', '') + '-int.tar'
    fn1 = fn.replace('.tar', '') + '-ext.tar'

    if not any(os.path.exists(x) for x in [fn, fn0, fn1]):
      ui.fatal('bundle file not found')

    ui.info(f'copying in: {fn} -> {target}')

    import time
    from subprocess import CalledProcessError
    from trueseeing.core.tools import invoke_passthru
    from trueseeing.core.android.device import AndroidDevice
    from trueseeing.core.android.tools import toolchains

    at = time.time()
    dev = AndroidDevice()

    if not await dev.is_fullbackup_available():
      ui.warn('full backup feature is not available')
    else:
      if not os.path.exists(fn):
        ui.warn(f'data not found, trying the next method: {fn}')
      else:
        try:
          try:
            with toolchains() as tc:
              await invoke_passthru('java -jar {abe} pack-kk {fn} {fn}.ab 1'.format(
                abe=tc['abe'],
                fn=fn,
              ))
          except CalledProcessError:
            ui.warn('pack failed; trying the next method')
          else:
            ui.success('pack success')
            ui.info('initiating a restore on device; give "1" as the password if asked')
            await dev.invoke_adb_passthru(f'restore {fn}.ab')
            ui.success(f'copied in: {fn}')
            success = True
        finally:
          try:
            os.remove(f'{fn}.ab')
          except FileNotFoundError:
            pass

    if not success:
      if not await dev.is_package_debuggable(target):
        ui.warn('target is not debuggable')
      else:
        ui.info('target seems debuggable; trying injection with debug interface')

        ui.info('copying internal storage')
        if not os.path.exists(fn0):
          ui.warn('data not found: {fn0}')
        else:
          tfn0 = self._generate_tempfilename_for_device()
          await dev.invoke_adb_passthru(f'push {fn0} {tfn0}')
          await dev.invoke_adb_passthru(f'shell "run-as {target} tar -xv < {tfn0}; rm -f {tfn0}"')
          ui.success(f'copied in: {fn}')
          success = True

        ui.info('copying external storage')
        if not os.path.exists(fn1):
          ui.warn('data not found: {fn1}')
        else:
          tfn1 = self._generate_tempfilename_for_device()
          await dev.invoke_adb_passthru(f'push {fn1} {tfn1}')
          await dev.invoke_adb_passthru(f'shell "cd /storage/emulated/0/Android/ && tar -xv < {tfn1}; rm -f {tfn1}"')
          ui.success(f'copied in: {fn1}')
          success = True

        success = True

    if success:
      ui.success('done ({t:.02f} sec.)'.format(t=(time.time() - at)))
    else:
      ui.failure('copyin failed')

  def _generate_tempfilename_for_device(self, dir: Optional[str] = None) -> str:
    import random
    return (f'{dir}/' if dir is not None else '/data/local/tmp/') + ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=16))

  def _parsed_manifest(self, blob: bytes) -> Any:
    import lxml.etree as ET
    return ET.fromstring(blob, parser=ET.XMLParser(recover=True))

  def _manifest_as_xml(self, manifest: Any) -> bytes:
    import lxml.etree as ET
    assert manifest is not None
    return ET.tostring(manifest) # type: ignore[no-any-return]
