# -*- coding: utf-8 -*-
# This file is part of the Smali Emulator.
#
# Copyright(c) 2016 Simone 'evilsocket' Margaritelli
# evilsocket@gmail.com
# http://www.evilsocket.net
#
# This file may be licensed under the terms of of the
# GNU General Public License Version 3 (the ``GPL'').
#
# Software distributed under the License is distributed
# on an ``AS IS'' basis, WITHOUT WARRANTY OF ANY KIND, either
# express or implied. See the GPL for the specific language
# governing rights and limitations.
#
# You should have received a copy of the GPL along with this
# program. If not, go to http://www.gnu.org/licenses/gpl.html
# or write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
from smaliemu.opcodes import *
from smaliemu.vm import VM
from smaliemu.source import Source
from smaliemu.snippet import Snippet
from smaliemu.preprocessors import TryCatchPreprocessor
from smaliemu.preprocessors import PackedSwitchPreprocessor
from smaliemu.preprocessors import ArrayDataPreprocessor
from smaliemu.exception import UnsupportedOpcodeError

import sys
import time


# Holds some statistics.
class Stats(object):
    def __init__(self, vm):
        self.opcodes = len(vm.opcodes)
        self.preproc = 0
        self.execution = 0
        self.steps = 0

    def __repr__(self):
        s = "opcode handlers    : %d\n" % self.opcodes
        s += "preprocessing time : %d ms\n" % self.preproc
        s += "execution time     : %d ms\n" % self.execution
        s += "execution steps    : %d\n" % self.steps
        return s


# The main emulator class.
class Emulator(object):
    def __init__(self):
        # Instance of the virtual machine.
        self.vm = None
        # Instance of the source file.
        self.source = None
        self.snippet = None
        # Instance of the statistics object.
        self.stats = None
        # Code preprocessors.
        self.preprocessors = [
            TryCatchPreprocessor, PackedSwitchPreprocessor,
            ArrayDataPreprocessor
        ]
        # Opcodes handlers.
        self.opcodes = []

        # Dynamically load opcode handlers.
        # TODO: Implement missing opcodes.
        for entry in dir(sys.modules['smaliemu.opcodes']):
            if entry.startswith('op_'):
                self.opcodes.append(globals()[entry]())

    def __preprocess(self, smali_code):
        """
        Start the preprocessing phase which will save all the labels and their
        line index for fast lookups while jumping and will pre parse all the
        try/catch directives.
        """
        next_line = None
        smali_code.lines = list(map(str.strip, smali_code.lines))
        for index, line in enumerate(smali_code.lines):
            # we're inside a block which was already processed
            if next_line is not None and index <= next_line:
                next_line = None if index == next_line else next_line
                continue
            elif line == '':  # skip empty lines
                continue
            elif line[0] == ':':  # we've found something to preprocess
                # loop each preprocessors and search for the one responsible
                # to parse this line
                processed = False
                for preproc in self.preprocessors:
                    if preproc.check(line):
                        next_line = preproc.process(self.vm, line, index,
                                                    smali_code.lines)
                        processed = True

                # no preprocessor found, this is a normal label
                if processed is False:
                    self.vm.labels[line] = index

    def __parse_line(self, line):
        # Search for appropriate parser.
        for parser in self.opcodes:
            if parser.parse(line, self.vm):
                return True

        return False

    @staticmethod
    def __should_skip_line(line):
        """
        Helper method used to determine if a line must be skipped/ignored.
        :param line: The line to check.
        :return: True if the line can be ignored, otherwise False.
        """
        return line == "" or line[0] == '#' or line[0] == ':' or line[0] == '.'

    def fatal(self, message):
        """
        Display an error message, the current line being executed and quit.
        :param message: The error message to display.
        """
        print()
        print("-------------------------")
        print("Fatal error on line %03d:\n" % self.vm.pc)
        if self.source:
            print("  %03d %s" % (self.vm.pc, self.source[self.vm.pc - 1]))
        elif self.snippet:
            print("  %03d %s" % (self.vm.pc, self.snippet[self.vm.pc - 1]))

        print("\n%s" % message)

    def get_error_line(self):
        if self.source:
            return "  %03d %s" % (self.vm.pc, self.source[self.vm.pc - 1])
        elif self.snippet:
            return "  %03d %s" % (self.vm.pc, self.snippet[self.vm.pc - 1])

    def run(self, filename, args={}, trace=False):
        """
        Load a smali file and start emulating it.
        :param filename: The path of the file to load and emulate.
        :param args: A dictionary of optional initialization variables for the
                     VM, mostly used for arguments.
        :param trace: If true every opcode being executed will be printed.
        :return: The return value of the emulated method or None if no return-*
                 opcode was executed.
        """
        OpCode.trace = trace
        self.source = Source(filename)
        self.vm = VM(self)
        self.stats = Stats(self)

        self.vm.variables.update(args)

        s = time.time() * 1000
        # Preprocess labels and try/catch blocks for fast lookup.
        self.__preprocess(self.source)
        e = time.time() * 1000
        self.stats.preproc = e - s

        s = time.time() * 1000
        # Loop each line and emulate.
        while self.vm.stop is False and self.source.has_line(self.vm.pc):
            self.stats.steps += 1
            line = self.source[self.vm.pc]
            self.vm.pc += 1

            if self.__should_skip_line(line):
                continue

            elif self.__parse_line(line) is False:
                raise UnsupportedOpcodeError(self.get_error_line())

        e = time.time() * 1000
        self.stats.execution = e - s

        return self.vm.return_v

    def call(self, snippet, trace=False, thrown=True):
        """
        Load a smali file and start emulating it.
        :param snippet: codes snippet in smali format
        :param trace: If true every opcode being executed will be printed.
        :param thrown: If error, emu will raise exception.
        :return: The return value of the emulated method or None if no return-*
                 opcode was executed.
        """
        OpCode.trace = trace
        self.snippet = Snippet(snippet)
        self.vm = VM(self)
        self.stats = Stats(self)
        self.vm.return_v = None
        self.vm.thrown = thrown

        self.vm.variables.copy()

        s = time.time() * 1000
        # Preprocess labels and try/catch blocks for fast lookup.
        self.__preprocess(self.snippet)
        e = time.time() * 1000
        self.stats.preproc = e - s

        s = time.time() * 1000
        # Loop each line and emulate.
        while self.vm.stop is False and self.snippet.has_line(self.vm.pc):
            self.stats.steps += 1
            line = self.snippet[self.vm.pc]
            self.vm.pc += 1

            if self.__should_skip_line(line):
                continue

            elif self.__parse_line(line) is False and thrown:
                raise UnsupportedOpcodeError(self.get_error_line())

        e = time.time() * 1000
        self.stats.execution = e - s

        return self.vm.return_v
