import inspect
from datetime import datetime
from ipaddress import IPv4Address

import pytest

from CybORG import CybORG
from CybORG.Shared.Actions import MSFAutoroute
from CybORG.Shared.Actions import SSHLoginExploit
from CybORG.Shared.Enums import SessionType, ProcessState, ProcessType, AppProtocol, Architecture, \
    OperatingSystemDistribution, OperatingSystemKernelVersion, OperatingSystemType
from CybORG.Tests.EphemeralPort import LinuxEphemeralPort

def test_ssh_brute_force_success():
    path = str(inspect.getfile(CybORG))
    path = path[:-10] + '/Shared/Scenarios/Scenario1.yaml'
    cyborg = CybORG(path, 'sim')
    agent = 'Red'
    initial_result = cyborg.get_observation(agent)
    # create ssh session on pretend pi host
    session = initial_result['Attacker']['Sessions'][0]['ID']
    k_ip_address = initial_result['Attacker']['Interface'][0]['IP Address']
    pp_ip_address = initial_result['Gateway']['Interface'][0]['IP Address']

    action = SSHLoginExploit(session=session, agent=agent, ip_address=pp_ip_address, port=22)
    expected_result = {"success": True,
                       str(pp_ip_address): {'Interface': [{'IP Address': pp_ip_address}],
                                            'Processes': [{'Connections': [{'Application Protocol': AppProtocol.SSH,
                                                                            'local_address': pp_ip_address,
                                                                            'local_port': 22,
                                                                            'Status': ProcessState.OPEN}],
                                                           'Process Type': ProcessType.SSH},
                                                          {'Connections': [{'local_address': pp_ip_address,
                                                                            'local_port': 22,
                                                                            'remote_address': k_ip_address,
                                                                            'remote_port': LinuxEphemeralPort()}]}],
                                            'Sessions': [{'Agent': 'Red',
                                                          'ID': 1,
                                                          'Type': SessionType.MSF_SHELL,
                                                          'Username': 'pi'}],
                                            'System info': {'Architecture': Architecture.x64,
                                                            'Hostname': 'Gateway',
                                                            'OSDistribution': OperatingSystemDistribution.UBUNTU,
                                                            'OSType': OperatingSystemType.LINUX},
                                            'User Info': [{'Password': 'raspberry',
                                                           'UID': 1001,
                                                           'Username': 'pi'}]
                                            },
                       str(k_ip_address): {
                           'Interface': [{'IP Address': k_ip_address}],
                           'Processes': [{'Connections': [{'remote_address': pp_ip_address,
                                                           'remote_port': 22,
                                                           'local_address': k_ip_address,
                                                           'local_port': LinuxEphemeralPort()}]}]
                       }}
    results = cyborg.step(agent, action, skip_valid_action_check=True)

    assert results.observation == expected_result

def test_sim_execute_inactive():
    path = str(inspect.getfile(CybORG))
    path = path[:-10] + '/Shared/Scenarios/Scenario1.yaml'
    cyborg = CybORG(path, 'sim')
    agent = 'Red'
    initial_result = cyborg.get_observation(agent)
    expected_observation = {"success": False}
    target = "10.0.0.2"
    state = cyborg.environment_controller.state

    session = state.sessions["Red"][0]

    session.active = False
    action = SSHLoginExploit(session=session.ident, agent="Red", ip_address=IPv4Address(target), port=22)
    observation = action.sim_execute(state)
    assert observation.get_dict() == expected_observation


def test_sim_execute_dead():
    path = str(inspect.getfile(CybORG))
    path = path[:-10] + '/Shared/Scenarios/Scenario1.yaml'
    cyborg = CybORG(path, 'sim')
    agent = 'Red'
    initial_result = cyborg.get_observation(agent)
    expected_observation = {"success": False}
    target = "10.0.0.2"
    state = cyborg.environment_controller.state

    session = state.sessions["Red"][0]

    state.remove_process(session.host, session.pid)
    action = SSHLoginExploit(session=session.ident, agent="Red", ip_address=IPv4Address(target), port=22)
    observation = action.sim_execute(state)
    assert observation.get_dict() == expected_observation

def test_ssh_brute_force_nacl_block():
    path = str(inspect.getfile(CybORG))
    path = path[:-10] + '/Shared/Scenarios/Scenario1.yaml'
    cyborg = CybORG(path, 'sim')
    agent = 'Red'
    initial_result = cyborg.get_observation(agent)
    # create ssh session on pretend pi host
    session = initial_result['Attacker']['Sessions'][0]['ID']
    k_ip_address = initial_result['Attacker']['Interface'][0]['IP Address']
    pp_ip_address = initial_result['Gateway']['Interface'][0]['IP Address']

    action = SSHLoginExploit(session=session, agent=agent, ip_address=pp_ip_address, port=23)
    expected_result = {"success": False}
    results = cyborg.step(agent, action)

    assert results.observation == expected_result
