# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
# You may obtain a copy of the License at: [package root]/LICENSE.txt
# Unless required by applicable law or agreed to in writing, software distributed under the License
#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and limitations under the License.

from __future__ import annotations

import logging
from typing import Any, Tuple, Iterator, Optional

import dnfile
from dncil.cil.body import CilMethodBody
from dncil.cil.error import MethodBodyFormatError
from dncil.clr.token import Token, StringToken, InvalidToken
from dncil.cil.body.reader import CilMethodBodyReaderBase

logger = logging.getLogger(__name__)

# key indexes to dotnet metadata tables
DOTNET_META_TABLES_BY_INDEX = {table.value: table.name for table in dnfile.enums.MetadataTables}


class DnfileMethodBodyReader(CilMethodBodyReaderBase):
    def __init__(self, pe: dnfile.dnPE, row: dnfile.mdtable.MethodDefRow):
        self.pe: dnfile.dnPE = pe
        self.offset: int = self.pe.get_offset_from_rva(row.Rva)

    def read(self, n: int) -> bytes:
        data: bytes = self.pe.get_data(self.pe.get_rva_from_offset(self.offset), n)
        self.offset += n
        return data

    def tell(self) -> int:
        return self.offset

    def seek(self, offset: int) -> int:
        self.offset = offset
        return self.offset


class DnClass(object):
    def __init__(self, token: int, namespace: str, classname: str):
        self.token: int = token
        self.namespace: str = namespace
        self.classname: str = classname

    def __hash__(self):
        return hash((self.token,))

    def __eq__(self, other):
        return self.token == other.token

    def __str__(self):
        return DnClass.format_name(self.namespace, self.classname)

    def __repr__(self):
        return str(self)

    @staticmethod
    def format_name(namespace: str, classname: str):
        name: str = classname
        if namespace:
            # like System.IO.File::OpenRead
            name = f"{namespace}.{name}"
        return name


class DnMethod(DnClass):
    def __init__(self, token: int, namespace: str, classname: str, methodname: str):
        super(DnMethod, self).__init__(token, namespace, classname)
        self.methodname: str = methodname

    def __str__(self):
        return DnMethod.format_name(self.namespace, self.classname, self.methodname)

    @staticmethod
    def format_name(namespace: str, classname: str, methodname: str):  # type: ignore
        # like File::OpenRead
        name: str = f"{classname}::{methodname}"
        if namespace:
            # like System.IO.File::OpenRead
            name = f"{namespace}.{name}"
        return name


class DnUnmanagedMethod:
    def __init__(self, token: int, modulename: str, methodname: str):
        self.token: int = token
        self.modulename: str = modulename
        self.methodname: str = methodname

    def __hash__(self):
        return hash((self.token,))

    def __eq__(self, other):
        return self.token == other.token

    def __str__(self):
        return DnUnmanagedMethod.format_name(self.modulename, self.methodname)

    def __repr__(self):
        return str(self)

    @staticmethod
    def format_name(modulename, methodname):
        return f"{modulename}.{methodname}"


def resolve_dotnet_token(pe: dnfile.dnPE, token: Token) -> Any:
    """map generic token to string or table row"""
    if isinstance(token, StringToken):
        user_string: Optional[str] = read_dotnet_user_string(pe, token)
        if user_string is None:
            return InvalidToken(token.value)
        return user_string

    table_name: str = DOTNET_META_TABLES_BY_INDEX.get(token.table, "")
    if not table_name:
        # table_index is not valid
        return InvalidToken(token.value)

    table: Any = getattr(pe.net.mdtables, table_name, None)
    if table is None:
        # table index is valid but table is not present
        return InvalidToken(token.value)

    try:
        return table.rows[token.rid - 1]
    except IndexError:
        # table index is valid but row index is not valid
        return InvalidToken(token.value)


def read_dotnet_method_body(pe: dnfile.dnPE, row: dnfile.mdtable.MethodDefRow) -> Optional[CilMethodBody]:
    """read dotnet method body"""
    try:
        return CilMethodBody(DnfileMethodBodyReader(pe, row))
    except MethodBodyFormatError as e:
        logger.warn("failed to parse managed method body @ 0x%08x (%s)" % (row.Rva, e))
        return None


def read_dotnet_user_string(pe: dnfile.dnPE, token: StringToken) -> Optional[str]:
    """read user string from #US stream"""
    try:
        user_string: Optional[dnfile.stream.UserString] = pe.net.user_strings.get_us(token.rid)
    except UnicodeDecodeError as e:
        logger.warn("failed to decode #US stream index 0x%08x (%s)" % (token.rid, e))
        return None

    if user_string is None:
        return None

    return user_string.value


def get_dotnet_managed_imports(pe: dnfile.dnPE) -> Iterator[DnMethod]:
    """get managed imports from MemberRef table

    see https://www.ntcore.com/files/dotnetformat.htm

    10 - MemberRef Table
        Each row represents an imported method
            Class (index into the TypeRef, ModuleRef, MethodDef, TypeSpec or TypeDef tables)
            Name (index into String heap)
    01 - TypeRef Table
        Each row represents an imported class, its namespace and the assembly which contains it
            TypeName (index into String heap)
            TypeNamespace (index into String heap)
    """
    for (rid, row) in enumerate(iter_dotnet_table(pe, "MemberRef")):
        if not isinstance(row.Class.row, dnfile.mdtable.TypeRefRow):
            continue

        token: int = calculate_dotnet_token_value(pe.net.mdtables.MemberRef.number, rid + 1)
        yield DnMethod(token, row.Class.row.TypeNamespace, row.Class.row.TypeName, row.Name)


def get_dotnet_managed_methods(pe: dnfile.dnPE) -> Iterator[DnMethod]:
    """get managed method names from TypeDef table

    see https://www.ntcore.com/files/dotnetformat.htm

    02 - TypeDef Table
        Each row represents a class in the current assembly.
            TypeName (index into String heap)
            TypeNamespace (index into String heap)
            MethodList (index into MethodDef table; it marks the first of a continguous run of Methods owned by this Type)
    """
    for row in iter_dotnet_table(pe, "TypeDef"):
        for index in row.MethodList:
            token = calculate_dotnet_token_value(index.table.number, index.row_index)
            yield DnMethod(token, row.TypeNamespace, row.TypeName, index.row.Name)


def get_dotnet_managed_method_bodies(pe: dnfile.dnPE) -> Iterator[Tuple[int, CilMethodBody]]:
    """get managed methods from MethodDef table"""
    if not hasattr(pe.net.mdtables, "MethodDef"):
        return

    for (rid, row) in enumerate(pe.net.mdtables.MethodDef):
        if not row.ImplFlags.miIL or any((row.Flags.mdAbstract, row.Flags.mdPinvokeImpl)):
            # skip methods that do not have a method body
            continue

        body: Optional[CilMethodBody] = read_dotnet_method_body(pe, row)
        if body is None:
            continue

        token: int = calculate_dotnet_token_value(dnfile.enums.MetadataTables.MethodDef.value, rid + 1)
        yield token, body


def get_dotnet_unmanaged_imports(pe: dnfile.dnPE) -> Iterator[DnUnmanagedMethod]:
    """get unmanaged imports from ImplMap table

    see https://www.ntcore.com/files/dotnetformat.htm

    28 - ImplMap Table
        ImplMap table holds information about unmanaged methods that can be reached from managed code, using PInvoke dispatch
            MemberForwarded (index into the Field or MethodDef table; more precisely, a MemberForwarded coded index)
            ImportName (index into the String heap)
            ImportScope (index into the ModuleRef table)
    """
    for row in iter_dotnet_table(pe, "ImplMap"):
        modulename: str = row.ImportScope.row.Name
        methodname: str = row.ImportName

        # ECMA says "Each row of the ImplMap table associates a row in the MethodDef table (MemberForwarded) with the
        # name of a routine (ImportName) in some unmanaged DLL (ImportScope)"; so we calculate and map the MemberForwarded
        # MethodDef table token to help us later record native import method calls made from CIL
        token: int = calculate_dotnet_token_value(row.MemberForwarded.table.number, row.MemberForwarded.row_index)

        # like Kernel32.dll
        if modulename and "." in modulename:
            modulename = modulename.split(".")[0]

        # like kernel32.CreateFileA
        yield DnUnmanagedMethod(token, modulename, methodname)


def calculate_dotnet_token_value(table: int, rid: int) -> int:
    return ((table & 0xFF) << Token.TABLE_SHIFT) | (rid & Token.RID_MASK)


def is_dotnet_table_valid(pe: dnfile.dnPE, table_name: str) -> bool:
    return bool(getattr(pe.net.mdtables, table_name, None))


def is_dotnet_mixed_mode(pe: dnfile.dnPE) -> bool:
    return not bool(pe.net.Flags.CLR_ILONLY)


def iter_dotnet_table(pe: dnfile.dnPE, name: str) -> Iterator[Any]:
    if not is_dotnet_table_valid(pe, name):
        return
    for row in getattr(pe.net.mdtables, name):
        yield row
