# Copyright (C) 2021 Mandiant, Inc. All Rights Reserved.

import json
import datetime
from enum import Enum
from typing import Dict, List
from dataclasses import field

# we use pydantic for dataclasses so that we can
# easily load and validate JSON reports.
#
# pydantic checks all the JSON fields look as they should
# while using the nice and familiar dataclass syntax.
#
# really, you should just pretend we're using stock dataclasses.
from pydantic.dataclasses import dataclass
from pydantic.error_wrappers import ValidationError

import floss.logging_
from floss.render import Verbosity
from floss.version import __version__
from floss.render.sanitize import sanitize

logger = floss.logging_.getLogger(__name__)


class InvalidResultsFile(Exception):
    pass


class InvalidLoadConfig(Exception):
    pass


class StringEncoding(str, Enum):
    ASCII = "ASCII"
    UTF16LE = "UTF-16LE"


@dataclass(frozen=True)
class StackString:
    """
    here's what the following members represent:
        
        
        [smaller addresses]
       
        +---------------+  <- stack_pointer (top of stack)
        |               | \
        +---------------+  | offset
        |               | /
        +---------------+
        | "abc"         | \
        +---------------+  |
        |               |  |
        +---------------+  | frame_offset
        |               |  |
        +---------------+  |
        |               | /
        +---------------+  <- original_stack_pointer (bottom of stack, probably bp)
        
        [bigger addresses]


    Attributes:
      function: the address of the function from which the stackstring was extracted
      string: the extracted string
      program_counter: the program counter at the moment the string was extracted
      stack_pointer: the stack counter at the moment the string was extracted
      original_stack_pointer: the initial stack counter when the function was entered
      offset: the offset into the stack from at which the stack string was found
      frame_offset: the offset from the function frame at which the stack string was found
    """

    function: int
    string: str
    encoding: StringEncoding
    program_counter: int
    stack_pointer: int
    original_stack_pointer: int
    offset: int
    frame_offset: int


class TightString(StackString):
    pass


class AddressType(str, Enum):
    STACK = "STACK"
    GLOBAL = "GLOBAL"
    HEAP = "HEAP"


@dataclass(frozen=True)
class DecodedString:
    """
    A decoding string and details about where it was found.

    Attributes:
        address: address of the string in memory
        address_type: type of the address of the string in memory
        string: the decoded string
        encoding: the string encoding, like ASCII or unicode
        decoded_at: the address at which the decoding routine is called
        decoding_routine: the address of the decoding routine
    """

    address: int
    address_type: AddressType
    string: str
    encoding: StringEncoding
    decoded_at: int
    decoding_routine: int


@dataclass(frozen=True)
class StaticString:
    """
    A string extracted from the raw bytes of the input.

    Attributes:
        string: the string
        offset: the offset into the input where the string is found
        encoding: the string encoding, like ASCII or unicode
    """

    string: str
    offset: int
    encoding: StringEncoding


@dataclass
class Runtime:
    start_date: datetime.datetime = datetime.datetime.now()
    total: float = 0
    vivisect: float = 0
    find_features: float = 0
    static_strings: float = 0
    stack_strings: float = 0
    decoded_strings: float = 0
    tight_strings: float = 0


@dataclass
class Functions:
    discovered: int = 0
    library: int = 0
    analyzed_stack_strings: int = 0
    analyzed_tight_strings: int = 0
    analyzed_decoded_strings: int = 0
    decoding_function_scores: Dict[int, float] = field(default_factory=dict)


@dataclass
class Analysis:
    enable_static_strings: bool = True
    enable_stack_strings: bool = True
    enable_tight_strings: bool = True
    enable_decoded_strings: bool = True
    functions: Functions = field(default_factory=Functions)


STRING_TYPE_FIELDS = set([field for field in Analysis.__annotations__ if field.startswith("enable_")])


@dataclass
class Metadata:
    file_path: str
    version: str = __version__
    imagebase: int = 0
    min_length: int = 0
    runtime: Runtime = field(default_factory=Runtime)


@dataclass
class Strings:
    stack_strings: List[StackString] = field(default_factory=list)
    tight_strings: List[TightString] = field(default_factory=list)
    decoded_strings: List[DecodedString] = field(default_factory=list)
    static_strings: List[StaticString] = field(default_factory=list)


@dataclass
class ResultDocument:
    metadata: Metadata
    analysis: Analysis = field(default_factory=Analysis)
    strings: Strings = field(default_factory=Strings)

    @classmethod
    def parse_file(cls, path):
        return cls.__pydantic_model__.parse_file(path)


def log_result(decoded_string, verbosity):
    string = sanitize(decoded_string.string)
    if verbosity < Verbosity.VERBOSE:
        logger.info("%s", string)
    else:
        if type(decoded_string) == DecodedString:
            logger.info(
                "%s [%s] decoded by 0x%x called at 0x%x",
                string,
                decoded_string.encoding,
                decoded_string.decoding_routine,
                decoded_string.decoded_at,
            )
        elif type(decoded_string) in (StackString, TightString):
            logger.info(
                "%s [%s] in 0x%x at address 0x%x",
                string,
                decoded_string.encoding,
                decoded_string.function,
                decoded_string.program_counter,
            )
        else:
            ValueError("unknown decoded or extracted string type: %s", type(decoded_string))


def load(sample: str, analysis: Analysis, functions: List[int], min_length: int) -> ResultDocument:
    logger.debug("loading results document: %s", sample)
    results = read(sample)
    results.metadata.file_path = f"{sample}\n{results.metadata.file_path}"
    check_set_string_types(results, analysis)
    if functions:
        filter_functions(results, functions)
    if min_length:
        filter_string_len(results, min_length)
        results.metadata.min_length = min_length
    return results


def read(sample: str) -> ResultDocument:
    try:
        with open(sample, "rb") as f:
            results = json.loads(f.read().decode("utf-8"))
    except (json.decoder.JSONDecodeError, UnicodeDecodeError) as e:
        raise InvalidResultsFile(f"{e}")

    try:
        results = ResultDocument(**results)
    except (TypeError, ValidationError) as e:
        raise InvalidResultsFile(f"{sample} is not a valid FLOSS result document: {e}")

    return results


def check_set_string_types(results: ResultDocument, wanted_analysis: Analysis) -> None:
    for string_type in STRING_TYPE_FIELDS:
        if getattr(wanted_analysis, string_type) and not getattr(results.analysis, string_type):
            logger.warning(f"{string_type} not in loaded data, use --only/--no to enable/disable type(s)")
        setattr(results.analysis, string_type, getattr(wanted_analysis, string_type))


def filter_functions(results: ResultDocument, functions: List[int]) -> None:
    filtered_scores = dict()
    for fva in functions:
        try:
            filtered_scores[fva] = results.analysis.functions.decoding_function_scores[fva]
        except KeyError:
            raise InvalidLoadConfig(f"function 0x{fva:x} not found in loaded data")
    results.analysis.functions.decoding_function_scores = filtered_scores

    results.strings.stack_strings = list(filter(lambda f: f.function in functions, results.strings.stack_strings))
    results.strings.tight_strings = list(filter(lambda f: f.function in functions, results.strings.tight_strings))
    results.strings.decoded_strings = list(
        filter(lambda f: f.decoding_routine in functions, results.strings.decoded_strings)
    )

    results.analysis.functions.analyzed_stack_strings = len(results.strings.stack_strings)
    results.analysis.functions.analyzed_tight_strings = len(results.strings.tight_strings)
    results.analysis.functions.analyzed_decoded_strings = len(results.strings.decoded_strings)


def filter_string_len(results: ResultDocument, min_length: int) -> None:
    results.strings.static_strings = list(filter(lambda s: len(s.string) >= min_length, results.strings.static_strings))
    results.strings.stack_strings = list(filter(lambda s: len(s.string) >= min_length, results.strings.stack_strings))
    results.strings.tight_strings = list(filter(lambda s: len(s.string) >= min_length, results.strings.tight_strings))
    results.strings.decoded_strings = list(
        filter(lambda s: len(s.string) >= min_length, results.strings.decoded_strings)
    )
