# Copyright DST Group. Licensed under the MIT license.
from ipaddress import IPv4Address

# For use after TomcatCredentialScanner most likely
# use exploit/multi/http/tomcat_mgr_upload
# set RHOSTS target, set RPORT target_port, set HttpPassword password, set HttpUsername user, set target 0
from CybORG.Shared.Actions.MSFActionsFolder.RemoteCodeExecutionFolder.RemoteCodeExecution import RemoteCodeExecution
from CybORG.Shared.Enums import InterfaceType, ProcessType, ProcessVersion, AppProtocol, SessionType
from CybORG.Shared.Observation import Observation


class TomcatExploit(RemoteCodeExecution):
    def __init__(self, session: int, agent: str, ip_address: IPv4Address, port: int, username: str,
                 password: str):
        super().__init__(session, agent)
        self.target = ip_address
        self.target_port = port
        self.user = username
        self.password = password

    def sim_execute(self, state):
        obs = Observation()
        obs.set_success(False)
        if self.session not in state.sessions[self.agent]:
            return obs
        from_host = state.sessions['Red'][self.session].host

        from_session = state.sessions['Red'][self.session]
        if from_session.session_type != SessionType.MSF_SERVER or not from_session.active:
            return obs

        from_session, from_interface = self.get_local_source_interface(local_session=from_session, remote_address=self.target)

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

        if self.target == IPv4Address("127.0.0.1"):
            target_host = from_host
        else:
            target_host = state.hosts[state.ip_addresses[self.target]]

        target_proc = None
        for process in target_host.processes:
            for conn in process.connections:
                if self.target_port == conn['local_port']:
                    if conn['local_address'] == IPv4Address("0.0.0.0") or (
                            conn['local_address'] == IPv4Address("127.0.0.1") and self.target == IPv4Address(
                            "127.0.0.1")) or conn['local_address'] == self.target:
                        target_proc = process
                    break

        if target_proc is None:
            obs.set_success(False)
            return obs
        else:
            if target_proc.process_type == ProcessType.WEBSERVER:
                if target_proc.version == ProcessVersion.APACHE_TOMCAT:
                    for conn in target_proc.connections:
                        if conn["local_port"] == self.target_port:
                            if conn['Application Protocol'] == AppProtocol.HTTP:
                                if self.user == "tomcat" and self.password == "tomcat":
                                    obs.set_success(True)
                                    obs.add_process(hostid="hostid1", local_address=self.target,
                                                    local_port=self.target_port,
                                                    status="open")
                                    obs.add_user_info(hostid="hostid1", username="tomcat", password="tomcat",
                                                      group_name="web server users")
                                    # TODO: add session that spawns from target_proc
                                    session = state.add_session(host=target_host.hostname, agent=self.agent,
                                                                user=target_proc.user.username,
                                                                session_type='Meterpreter', parent=from_session)
                                    port = target_host.get_ephemeral_port()
                                    session.process.user = session.user
                                    session.process.name = "java"
                                    session.process.path = "/usr/bin/"  # Should acutally be "/usr/lib/jvm/java-1.5.0-gcj-4.2-1.5.0.0/jre/bin/" ?
                                    session.process.ppid = 1
                                    new_connection = {
                                        'local_port': port,
                                        'local_address': self.target,
                                        'remote_port': 4444,
                                        'remote_address': from_interface.ip_address,
                                        'Application Protocol': "tcp"
                                    }
                                    session.process.connections.append(new_connection)

                                    remote_port = {
                                        'local_port': 4444,
                                        "Application Protocol": "tcp",
                                        "local_address": from_interface.ip_address,
                                        "remote_address": self.target,
                                        "remote_port": port
                                    }
                                    from_session.process.connections.append(remote_port)

                                    obs.add_session_info(hostid='hostid1', timeout=0, session_id=session.ident,
                                                         session_type=session.session_type, agent=self.agent)
                                    obs.add_process(hostid='hostid1', local_address=self.target,
                                                    remote_address=from_interface.ip_address, remote_port=4444,
                                                    local_port=port)
                                    obs.add_process(hostid='hostid0', local_address=from_interface.ip_address,
                                                    remote_address=self.target, local_port=4444, remote_port=port)
        return obs

    def emu_execute(self, session_handler) -> Observation:
        obs = Observation()
        from CybORG.Emulator.Session import MSFSessionHandler
        if type(session_handler) is not MSFSessionHandler:
            obs.set_success(False)
            return obs
        output = session_handler.execute_exploit(exploit_name='multi/http/tomcat_mgr_upload', payload='java/meterpreter/bind_tcp',
                                         opts={'RHOSTS': str(self.target), "RPORT": self.target_port, "HttpPassword": self.password, "HttpUsername": self.user})
        obs.add_raw_obs(output)
        obs.set_success(False)
        return obs