## 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.
"""
Pertaining to actions that establish a new session on a host by exploiting
a process on that host
"""
# pylint: disable=invalid-name
from ipaddress import IPv4Address
from random import random
from typing import Tuple, Optional
from abc import abstractmethod
from csle_cyborg.shared.observation import Observation
from csle_cyborg.shared.actions.concrete_actions.concrete_action import ConcreteAction
from csle_cyborg.shared.actions.msf_actions_folder.msf_action import lo
from csle_cyborg.simulator.process import ProcessType, DecoyType
from csle_cyborg.shared.enums import OperatingSystemType
from csle_cyborg.simulator.host import Host
from csle_cyborg.simulator.process import Process
from csle_cyborg.simulator.session import Session
from csle_cyborg.simulator.state import State


class ExploitAction(ConcreteAction):
    """
    Base class for actions that establish a new session on a host by exploiting
    a process on that host
    """

    def __init__(self, session: int, agent: str, ip_address: IPv4Address, target_session: int):
        super().__init__(session, agent)
        self.state = None
        self.ip_address = ip_address
        self.target_session = target_session
        self.detection_rate = 0.95

    def _get_originating_ip(self, state, from_host, port) -> \
            Optional[Tuple[Host, Optional[IPv4Address]]]:
        originating_ip_address = None
        if self.ip_address == lo:
            target_host: Host = from_host
            originating_ip_address = self.ip_address
        else:
            target_host: Host = state.hosts[state.ip_addresses[self.ip_address]]
            ports = self.check_routable(
                [state.subnets[i.subnet] for i in from_host.interfaces
                 if i.ip_address != lo],
                [s for s in state.subnets.values() if self.ip_address in s.cidr])
            if ports is None or (port not in ports and 'all' not in ports):
                return None
            from_subnet, _to_subnet = ports[port] if port in ports else ports['all']
            # calculate the originating ip address
            for i in from_host.interfaces:
                if i.ip_address != lo and i.subnet == from_subnet:
                    originating_ip_address = i.ip_address

        return target_host, originating_ip_address

    def sim_exploit(self, state: State, port: int, process_type: str) -> Observation:
        """Attempts to exploit host, establishing a session on it"""
        self.state = state
        obs = Observation()
        if self.session not in state.sessions[self.agent]:
            obs.set_success(False)
            return obs
        session = state.sessions[self.agent][self.session]
        from_host = state.hosts[session.host]

        if not session.active:
            obs.set_success(False)
            return obs

        # determine which ports can communicate between which subnets
        res = self._get_originating_ip(state, from_host, port)
        if res is None:
            obs.set_success(False)
            return obs

        target_host, originating_ip_address = res

        # find out if smb or smtp or femitter is running
        vuln_proc = None
        for proc in target_host.processes:
            if proc.process_type == ProcessType.parse_string(process_type):
                for conn in proc.connections:
                    if 'local_port' in conn and conn['local_port'] == port:
                        vuln_proc = proc
                        target_host.events['NetworkConnections'].append(
                            {'local_address': self.ip_address,
                             'local_port': conn["local_port"],
                             'remote_address': originating_ip_address,
                             'remote_port': target_host.get_ephemeral_port()})
                        break
                if vuln_proc is not None:
                    break

        if vuln_proc is None:
            obs.set_success(False)
            return obs

        if self.test_exploit_works(target_host, vuln_proc):
            if bool(vuln_proc.decoy_type & DecoyType.EXPLOIT):
                obs.add_interface_info(ip_address=str(self.ip_address))
                obs.set_success(False)
                target_host.events['NetworkConnections'].append({'local_address': self.ip_address,
                                                                 'local_port': port,
                                                                 'remote_address': originating_ip_address})

            elif bool(vuln_proc.decoy_type & DecoyType.SANDBOXING_EXPLOIT):
                obs = self.__add_reverse_session(state=state, username=vuln_proc.user,
                                                 from_ip_address=originating_ip_address, target_host=target_host,
                                                 session=session, is_escalate_sandbox=True)

            else:
                obs = self.__add_reverse_session(state=state, username=vuln_proc.user,
                                                 from_ip_address=originating_ip_address, target_host=target_host,
                                                 session=session)
        else:
            obs.add_interface_info(ip_address=str(self.ip_address))
            obs.set_success(False)

        obs.add_process(hostid=str(self.ip_address), local_address=self.ip_address,
                        local_port=port, status="open",
                        process_type=process_type)

        return obs

    @abstractmethod
    def test_exploit_works(self, target_host: Host,
                           vuln_process: Process) -> bool:
        """
        check if OS and process information is correct for exploit to work.
        Return True if it will work given the vulnerable process vuln_process
        False otherwise
        """
        raise NotImplementedError

    def __add_reverse_session(self, *, state: State, username: str,
                              from_ip_address: IPv4Address, target_host: Host,
                              session: Session, is_escalate_sandbox: bool = False):
        """
        called when exploit is successful, establishes a new session on that host
        """
        if target_host.os_type == OperatingSystemType.WINDOWS:
            ext = 'exe'
            path = 'C:\\temp\\'
        elif target_host.os_type == OperatingSystemType.LINUX:
            ext = 'sh'
            path = '/tmp/'
        else:
            return Observation(False)
        obs = Observation()

        local_port = target_host.get_ephemeral_port()
        new_connection = {"local_port": local_port,
                          "Application Protocol": "tcp",
                          "remote_address": from_ip_address,
                          "remote_port": 4444,
                          "local_address": self.ip_address
                          }

        file = target_host.add_file(f'cmd.{ext}', path, username, 7, density=0.9, signed=False)
        proc = target_host.add_process(name=file.name, ppid=1, path=file.path, user=username)
        proc.connections.append(new_connection)
        new_session = state.add_session(host=target_host.hostname, agent=self.agent,
                                        user=username, session_type="red_reverse_shell", parent=self.session,
                                        process=proc.pid, is_escalate_sandbox=is_escalate_sandbox)
        if random() < self.detection_rate:
            target_host.events['NetworkConnections'].append({'local_address': self.ip_address,
                                                             'local_port': local_port,
                                                             'remote_address': from_ip_address,
                                                             'remote_port': 4444,
                                                             'pid': proc.pid})
        remote_port = {"local_port": 4444,
                       "Application Protocol": "tcp",
                       "local_address": from_ip_address,
                       "remote_address": self.ip_address,
                       "remote_port": local_port
                       }

        state.hosts[session.host].get_process(session.pid).connections.append(remote_port)
        obs.add_process(hostid=str(from_ip_address),
                        local_address=from_ip_address,
                        remote_address=str(self.ip_address),
                        local_port=4444,
                        remote_port=local_port,
                        process_type='reverse_session_handler')
        obs.add_process(hostid=str(self.ip_address),
                        local_address=str(self.ip_address),
                        remote_address=from_ip_address,
                        local_port=local_port,
                        remote_port=4444,
                        process_type='reverse_session')
        obs.add_session_info(hostid=str(self.ip_address),
                             session_id=new_session.ident,
                             session_type=new_session.session_type,
                             agent=self.agent)
        obs.add_system_info(hostid=str(self.ip_address), hostname=target_host.hostname,
                            os_type=target_host.os_type)
        obs.set_success(True)
        return obs
