## The following code contains work of the United States Government and is not subject to domestic copyright protection under 17 USC § 105.
## Additionally, we waive copyright and related rights in the utilized code worldwide through the CC0 1.0 Universal public domain dedication.

import random
from ipaddress import IPv4Address
from abc import ABC, abstractmethod
from typing import Optional

from csle_cyborg.shared.observation import Observation
from csle_cyborg.shared.actions.action import Action
from csle_cyborg.shared.actions.concrete_actions.exploit_action import ExploitAction
from csle_cyborg.shared.actions.concrete_actions.blue_keep import BlueKeep
from csle_cyborg.shared.actions.concrete_actions.eternal_blue import EternalBlue
from csle_cyborg.shared.actions.concrete_actions.http_rfi import HTTPRFI
from csle_cyborg.shared.actions.concrete_actions.http_srfi import HTTPSRFI
from csle_cyborg.shared.actions.concrete_actions.ssh_brute_force import SSHBruteForce
from csle_cyborg.shared.actions.concrete_actions.ftp_directory_traversal import FTPDirectoryTraversal
from csle_cyborg.shared.actions.concrete_actions.sql_injection import SQLInjection
from csle_cyborg.shared.actions.concrete_actions.haraka_rce import HarakaRCE
from csle_cyborg.simulator.session import RedAbstractSession
from csle_cyborg.simulator.state import State

# pylint: disable=too-few-public-methods
class ExploitActionSelector(ABC):
    """
    Examines the target host and returns a selected applicable escalate action
    if any, as well as processes that are required to be genuine
    """
    # pylint: disable=missing-function-docstring
    @abstractmethod
    def get_exploit_action(self, *, state: State, session: int,
            agent: str, ip_address: IPv4Address) -> Optional[ExploitAction]:
        pass

class DefaultExploitActionSelector(ExploitActionSelector):
    """
    Chooses exploit based on available ports
    """
    def get_exploit_action(self, *, state: State, session: int,
            agent: str, ip_address: IPv4Address, priority=None) -> Optional[ExploitAction]:
        ports = state.sessions[agent][session].ports[ip_address]
        # Weights need to be different. Should be refactored to a list.
        exploit_options = {EternalBlue: 2.0 if 139 in ports else None,
                            BlueKeep: 1.0 if 3389 in ports else None,
                            HTTPRFI: 3.0 if 80 in ports else None,
                            HTTPSRFI: 4.0 if 443 in ports else None,
                            SSHBruteForce: 0.1 if 22 in ports else None,
                            SQLInjection: 5.0 if (3390 in ports) and (80 in ports or 443 in ports) \
                                    else None,
                            HarakaRCE: 6.0 if 25 in ports else None,
                            FTPDirectoryTraversal: 7.0 if 21 in ports else None,
                            }
        exploit_options = {a: x for a, x in exploit_options.items() if x is not None}
        if len(exploit_options) > 1:
            # use information to populate weights for which exploit to select
            odds_of_top_choice = 0.75
            top_choice = max(list(exploit_options.values()))
            other_choice = random.choices([x for x in list(exploit_options.values()) if x != top_choice])[0]
            # other_choice = random.choices([x for x in list(exploit_options.values())])[0]
            priority_value = exploit_options.get(priority)
            if priority_value is not None:
                selected_choice = priority_value
            elif random.random() < odds_of_top_choice:
                selected_choice = top_choice
            else:
                selected_choice = other_choice

            # sample the exploit to perform from the exploit weights
            sub_action = list(exploit_options.keys())[
                    list(exploit_options.values()).index(selected_choice)](session=session,
                                                                                            agent=agent,
                                                                                            ip_address=ip_address,
                                                                                            target_session=session)
            return sub_action
        elif len(exploit_options) == 1:
            return list(exploit_options.keys())[0](session=session,
                                                                                            agent=agent,
                                                                                            ip_address=ip_address,
                                                                                            target_session=session)
        else:
            return None


_default_exploit_action_selector = DefaultExploitActionSelector()

class ExploitRemoteService(Action):
    def __init__(self, session: int, agent: str, ip_address: IPv4Address, priority=None):
        super().__init__()
        self.ip_address = ip_address
        self.agent = agent
        self.session = session
        self.exploit_action_selector = _default_exploit_action_selector
        self.priority = priority

    def sim_execute(self, state: State) -> Observation:
        # find session inside or close to the target subnet
        if type(state.sessions[self.agent][self.session]) is not RedAbstractSession or self.ip_address not in state.sessions[self.agent][self.session].ports:
            return Observation(success=False)

        self.sub_action = self.exploit_action_selector.get_exploit_action(state=state,
                session=self.session, agent=self.agent, ip_address=self.ip_address, priority=self.priority)
        if self.sub_action is None:
            return Observation(success=False)
        obs = self.sub_action.sim_execute(state)
        if self.ip_address != IPv4Address('127.0.0.1') and obs.data['success'] == True:
            hostname = obs.data[str(self.ip_address)]["System info"]["Hostname"]
            os = obs.data[str(self.ip_address)]["System info"]["OSType"]
            state.sessions[self.agent][self.session].addos(hostname, os)
        return obs

    def __str__(self):
        return f"{self.__class__.__name__} {self.ip_address}"
