# Copyright 2020 Google LLC
#
# 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.


import codecs
from typing import Iterator, Optional

import idc
import idaapi

import capa.ida.helpers
from capa.features.address import Address, FileOffsetAddress, AbsoluteVirtualAddress
from capa.ida.plugin.qt_compat import QtCore, qt_get_item_flag_tristate


def info_to_name(display):
    """extract root value from display name

    e.g. function(my_function) => my_function
    """
    try:
        return display.split("(")[1].rstrip(")")
    except IndexError:
        return ""


def ea_to_hex(ea):
    """convert effective address (ea) to hex for display"""
    return f"{hex(ea)}"


class CapaExplorerDataItem:
    """store data for CapaExplorerDataModel"""

    def __init__(self, parent: Optional["CapaExplorerDataItem"], data: list[str], can_check=True):
        """initialize item"""
        self.pred = parent
        self._data = data
        self._children: list["CapaExplorerDataItem"] = []
        self._checked = False
        self._can_check = can_check

        # default state for item
        self.flags = QtCore.Qt.ItemIsEnabled | QtCore.Qt.ItemIsSelectable

        if self._can_check:
            self.flags = self.flags | QtCore.Qt.ItemIsUserCheckable | qt_get_item_flag_tristate()

        if self.pred:
            self.pred.appendChild(self)

    def setIsEditable(self, isEditable=False):
        """modify item editable flags

        @param isEditable: True, can edit, False cannot edit
        """
        if isEditable:
            self.flags |= QtCore.Qt.ItemIsEditable
        else:
            self.flags &= ~QtCore.Qt.ItemIsEditable

    def setChecked(self, checked):
        """set item as checked

        @param checked: True, item checked, False item not checked
        """
        self._checked = checked

    def canCheck(self):
        """ """
        return self._can_check

    def isChecked(self):
        """get item is checked"""
        return self._checked

    def appendChild(self, item: "CapaExplorerDataItem"):
        """add a new child to specified item

        @param item: CapaExplorerDataItem
        """
        self._children.append(item)

    def child(self, row: int) -> "CapaExplorerDataItem":
        """get child row

        @param row: row number
        """
        return self._children[row]

    def childCount(self) -> int:
        """get child count"""
        return len(self._children)

    def columnCount(self) -> int:
        """get column count"""
        return len(self._data)

    def data(self, column: int) -> Optional[str]:
        """get data at column

        @param: column number
        """
        try:
            return self._data[column]
        except IndexError:
            return None

    def parent(self) -> Optional["CapaExplorerDataItem"]:
        """get parent"""
        return self.pred

    def row(self) -> int:
        """get row location"""
        if self.pred:
            return self.pred._children.index(self)
        return 0

    def setData(self, column: int, value: str):
        """set data in column

        @param column: column number
        @value: value to set (assume str)
        """
        self._data[column] = value

    def children(self) -> Iterator["CapaExplorerDataItem"]:
        """yield children"""
        yield from self._children

    def removeChildren(self):
        """remove children"""
        del self._children[:]

    def __str__(self):
        """get string representation of columns

        used for copy-n-paste operations
        """
        return " ".join([data for data in self._data if data])

    @property
    def info(self):
        """return data stored in information column"""
        return self._data[0]

    @property
    def location(self) -> Optional[int]:
        """return data stored in location column"""
        try:
            # address stored as str, convert to int before return
            return int(self._data[1], 16)
        except ValueError:
            return None

    @property
    def details(self):
        """return data stored in details column"""
        return self._data[2]


class CapaExplorerRuleItem(CapaExplorerDataItem):
    """store data for rule result"""

    fmt = "%s (%d matches)"

    def __init__(
        self, parent: CapaExplorerDataItem, name: str, namespace: str, count: int, source: str, can_check=True
    ):
        """initialize item

        @param parent: parent node
        @param name: rule name
        @param namespace: rule namespace
        @param count: number of match for this rule
        @param source: rule source (tooltip)
        """
        display = self.fmt % (name, count) if count > 1 else name
        super().__init__(parent, [display, "", namespace], can_check)
        self._source = source

    @property
    def source(self):
        """return rule source to display (tooltip)"""
        return self._source


class CapaExplorerRuleMatchItem(CapaExplorerDataItem):
    """store data for rule match"""

    def __init__(self, parent: CapaExplorerDataItem, display: str, source=""):
        """initialize item

        @param parent: parent node
        @param display: text to display in UI
        @param source: rule match source to display (tooltip)
        """
        super().__init__(parent, [display, "", ""])
        self._source = source

    @property
    def source(self):
        """return rule contents for display"""
        return self._source


class CapaExplorerFunctionItem(CapaExplorerDataItem):
    """store data for function match"""

    fmt = "function(%s)"

    def __init__(self, parent: CapaExplorerDataItem, location: Address, can_check=True):
        """initialize item

        @param parent: parent node
        @param location: virtual address of function as seen by IDA
        """
        assert isinstance(location, AbsoluteVirtualAddress)
        ea = int(location)
        super().__init__(parent, [self.fmt % idaapi.get_name(ea), ea_to_hex(ea), ""], can_check)

    @property
    def info(self):
        """return function name"""
        info = super().info
        display = info_to_name(info)
        return display if display else info

    @info.setter
    def info(self, display):
        """set function name

        called when user changes function name in plugin UI

        @param display: new function name to display
        """
        self._data[0] = self.fmt % display


class CapaExplorerSubscopeItem(CapaExplorerDataItem):
    """store data for subscope match"""

    fmt = "subscope(%s)"

    def __init__(self, parent: CapaExplorerDataItem, scope):
        """initialize item

        @param parent: parent node
        @param scope: subscope name
        """
        super().__init__(parent, [self.fmt % scope, "", ""])


class CapaExplorerBlockItem(CapaExplorerDataItem):
    """store data for basic block match"""

    fmt = "basic block(loc_%08X)"

    def __init__(self, parent: CapaExplorerDataItem, location: Address):
        """initialize item

        @param parent: parent node
        @param location: virtual address of basic block as seen by IDA
        """
        assert isinstance(location, AbsoluteVirtualAddress)
        ea = int(location)
        super().__init__(parent, [self.fmt % ea, ea_to_hex(ea), ""])


class CapaExplorerInstructionItem(CapaExplorerBlockItem):
    """store data for instruction match"""

    fmt = "instruction(loc_%08X)"


class CapaExplorerDefaultItem(CapaExplorerDataItem):
    """store data for default match e.g. statement (and, or)"""

    def __init__(
        self, parent: CapaExplorerDataItem, display: str, details: str = "", location: Optional[Address] = None
    ):
        """initialize item

        @param parent: parent node
        @param display: text to display in UI
        @param details: text to display in details section of UI
        @param location: virtual address as seen by IDA
        """
        ea = None
        if location:
            assert isinstance(location, AbsoluteVirtualAddress)
            ea = int(location)

        super().__init__(parent, [display, ea_to_hex(ea) if ea is not None else "", details])


class CapaExplorerFeatureItem(CapaExplorerDataItem):
    """store data for feature match"""

    def __init__(
        self, parent: CapaExplorerDataItem, display: str, location: Optional[Address] = None, details: str = ""
    ):
        """initialize item

        @param parent: parent node
        @param display: text to display in UI
        @param details: text to display in details section of UI
        @param location: virtual address as seen by IDA
        """
        if location:
            assert isinstance(location, (AbsoluteVirtualAddress, FileOffsetAddress))
            ea = int(location)
            super().__init__(parent, [display, ea_to_hex(ea), details])
        else:
            super().__init__(parent, [display, "", details])


class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem):
    """store data for instruction match"""

    def __init__(self, parent: CapaExplorerDataItem, display: str, location: Address):
        """initialize item

        details section shows disassembly view for match

        @param parent: parent node
        @param display: text to display in UI
        @param location: virtual address as seen by IDA
        """
        assert isinstance(location, AbsoluteVirtualAddress)
        ea = int(location)
        details = capa.ida.helpers.get_disasm_line(ea)
        super().__init__(parent, display, location=location, details=details)
        self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)


class CapaExplorerByteViewItem(CapaExplorerFeatureItem):
    """store data for byte match"""

    def __init__(self, parent: CapaExplorerDataItem, display: str, location: Address):
        """initialize item

        details section shows byte preview for match

        @param parent: parent node
        @param display: text to display in UI
        @param location: virtual address as seen by IDA
        """
        assert isinstance(location, (AbsoluteVirtualAddress, FileOffsetAddress))
        ea = int(location)

        byte_snap = idaapi.get_bytes(ea, 32)

        details = ""
        if byte_snap:
            byte_snap = codecs.encode(byte_snap, "hex").upper()
            details = " ".join([byte_snap[i : i + 2].decode() for i in range(0, len(byte_snap), 2)])

        super().__init__(parent, display, location=location, details=details)
        self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)


class CapaExplorerStringViewItem(CapaExplorerFeatureItem):
    """store data for string match"""

    def __init__(self, parent: CapaExplorerDataItem, display: str, location: Address, value: str):
        """initialize item

        @param parent: parent node
        @param display: text to display in UI
        @param location: virtual address as seen by IDA
        """
        assert isinstance(location, (AbsoluteVirtualAddress, FileOffsetAddress))
        ea = int(location)

        super().__init__(parent, display, location=location, details=value)
        self.ida_highlight = idc.get_color(ea, idc.CIC_ITEM)
