Building your agents

In this section we will guide you through the process of creating an agent that you can use for the challenge. But before we dive into the agent's internals, let's step back a little and discuss where and how are the agents actually interacting with the rest of their world.

Bird's view of the challenge concepts and internals

AICA challenge is built on top of CYST, which is a multi-agent discrete-event simulation framework tailored for cybersecurity domain. This dictates many of the design choices, but the two most important concepts are:

• discrete events,
• message passing.

What these means in practice is that within challenge runs:

• everything happens in discrete time slices and any effect can be attributed to a specific time slice,
• multiple events and effects can happen in the same time slice and their order is not specified,
• every effect is communicated via messages that travel through the simulated infrastructure and there are no side channel communications.

In other words, agents operate in fully simulated time and interact with their environment only through sending and receiving messages. These messages - request and responses - carry agent's actions and reactions to these actions. There are however two things these messages do not carry:

• state information,
• rewards for their actions.

In this, CYST differs from the likes of OpenAI gym and subsequently, AICA challenge differs from other challenges, such as the CAGE Challenge. Our goal is to steer the community to creation of autonomous cybersecurity solutions that can operate in the real world, with minimal supervision. In such setting, they will not be receiving neat state updates and will have to infer themselves, if what they are doing is correct or not. Also, the state space dictates how agents can reason about the world around them. Therefore, state space construction is a crucial element of building an agent. Don't worry - the messages contain enough information to construct state spaces to your liking.

That being said, it's time to create your first agent...

Creating a new agent

The agents exist as separate Python packages that conform to CYST agent specification. However, to make it easy for you, it is best to use the challenge TUI to create the skeleton for you.

So, run the TUI...

  $ poetry run aica_challenge

... and navigate to the Init agent option:

You will be prompted to provide the agent's name. This should be Python package compliant (e.g., my-long-name). If you provide a non-compliant name, it will attempt to rename it to a compliant version.

If everything went ok, you should see the following screen prompting you to install the agent:

It is usually safe to let the system install the agent right away, as the installation happens in editable mode, so any changes you make to the agent later will be reflected in the challenge execution. However, if you decide to skip it, you can install it anytime later through the Install package option of Package/Agent management.

Now... what has the agent initialization actually done? It created an agent skeleton in the agents/ folder of your challenge directory. The tree looks like this:

└─ my-awesome-agent
   └─src
   | └─ my_awesome_agent
   | |  └─ main.py
   | |  └─ __init__.py
   | └─ __init__.py
   └─ pyproject.toml
   └─ README.md

The two important files there are:

• pyproject.toml for specifying dependencies and the whole package building
• main.py for the rest, which we will discuss in the following sections.

Agent's skeleton

If you open the main.py file, you will see the following skeleton:

import logging

from abc import ABC, abstractmethod
from typing import List, Tuple, Optional, Dict, Any, Union, Callable, Awaitable

from cyst.api.logic.action import Action
from cyst.api.logic.access import Authorization, AuthenticationToken
from cyst.api.environment.environment import EnvironmentMessaging
from cyst.api.environment.message import Request, Response, MessageType, Message
from cyst.api.environment.resources import EnvironmentResources
from cyst.api.network.session import Session
from cyst.api.host.service import ActiveService, ActiveServiceDescription, Service


class MyAgent(ActiveService):
    def __init__(self, msg: EnvironmentMessaging, res: EnvironmentResources, id: str, args: Optional[Dict[str, Any]]):
        pass

    async def run(self):
        pass

    async def process_message(self, message: Message) -> Tuple[bool, int]:
        pass


def create_agent(msg: EnvironmentMessaging, res: EnvironmentResources, id:str, args: Optional[Dict[str, Any]]) -> ActiveService:
    actor = MyAgent(msg, res, id, args)
    return actor


service_description = ActiveServiceDescription(
    "my-awesome-agent",
    "A description of an agent",
    create_agent
)

So, let's unpack it...

service_description = ActiveServiceDescription(
    "my-awesome-agent",
    "A description of an agent",
    create_agent
)

This is the entry point that is used by CYST to register the agent into the system. There is not really that much to change, aside maybe from the description in the second parameter. While technically not needed for you, it is something you will see if you list the installed agents. If you can't do descriptive names, this is your friend.

You can also change the name in the first parameter. Beware, however, that if you do that your previously configured challenge runs may stop working (more on that in the challenge running section).

def create_agent(msg: EnvironmentMessaging, res: EnvironmentResources, id:str, args: Optional[Dict[str, Any]]) -> ActiveService:
    actor = MyAgent(msg, res, id, args)
    return actor

This is something you really should not need to touch. It is a factory function for agent creation that is as simple as it gets. Better do any changes in the initializer.

from cyst.api... import

The whole AICA challenge is built on top of CYST simulation platform and an agent is just one of CYST's classes. Therefore, you may need to include other CYST declarations (even though, most what you need should be here already). Details can be found in the API reference or mostly discovered in the following text.

class MyAgent(ActiveService):
    def __init__(self, msg: EnvironmentMessaging, res: EnvironmentResources, id: str, args: Optional[Dict[str, Any]]):
        pass

    async def run(self):
        pass

    async def process_message(self, message: Message) -> Tuple[bool, int]:
        pass

Finally, here is the flesh and bones of your agent. As you can see, there is not much here, so we will give just a brief comment on the purpose of the functions and the rest will be explained in the following section, when we give you an introduction to creating an agent that does at least something.

In the __init__() function, your agent gets all the information and resources it needs to start. This is in fact the only place, where the agent receives any information aside from the message reception and occasional signalling from the environment.

The run() function is executed only once, when the agent is starting (not necessarily at the beginning of a challenge run). Here, the agent must make its first action, otherwise, it won't be doing anything, unless a message from another agent awakens it.

The process_message() function is where agent's interactions happen. Whenever an agent receives a message, this function is called and the agent has to react to it somehow, preferably with sending another message.

This is in fact the whole modus operandi of CYST agents. They send and receive messages and react on what is in the messages. Because the underlying CYST platform is multi-agent (and actually time-agnostic), many things can happen between agents' requests and the responses to them. An agent can even interleave action processing by doing many short actions, while long-term actions are unfolding. Therefore, you can't really expect an API like this:

observation, reward, done, info = env.step(action)

Everything is asynchronous, more complicated, but also more realistic...

Building a pseudorandom agent

In this section we will guide you through building a pseudorandom agent. This agent will be (given enough time) able to finish most of the simpler scenarios.

Agent initialization

Let's start with the __init__() function.

First we store all the parameters we get, as we may use them later:

def __init__(self, msg: EnvironmentMessaging, res: EnvironmentResources, id: str, args: Optional[Dict[str, Any]]):
   self._msg = msg
   self._res = res
   self._id = id
   self._args = args

EnvironmentMessaging and EnvironmentResources are described in great detail in the API reference, however, for now you really only need to know two things:

• EnvironmentMessaging lets you create and send messages to interact with the rest of the system.
• EnvironmentResources gives you access to actions, exploits and other useful stuff.

Action preparation

We will also store actions, so that we do not need to always go through resources and its limited interface.

   self._actions = res.action_store.get_prefixed("ac1")

The action_store enables you to retrieve (and also create) actions. Actions spaces are grouped by a namespace tag, which for the challenge is ac1. In the default installation, CYST also contains other action spaces, so we use this one to not have to deal with other more or less compatible actions.

What is in the self._actions after the previous line executes? As usual, the reference gives the exhaustive answer, but for you the most important thing is that you will have a list of Action objects that can be passed to messages and that will have some effect.

The most important attributes of the Action are:

• id: The identification of an action (or its name if you will). The id is a fully-qualified name including the namespace, e.g., "ac1:scan_host".
• description: A textual explanation of action purpose.
• parameters: Parameters (if any) the action requires. More on that later.
• exploit: An exploit (if needed) that should enable the action to successfully execute. Again, more on that later.

Testing the agent

Ok... now you have an agent with an __init__() function, and you probably want to test if everything works. The machinery behind running the whole simulation is not really straightforward, so, if you haven't checked running the challenge section, I will ask you to visit it and prepare a suitable run specification with these parameters:

• run name = test_run, so that you can refer to it from the testing script,
• scenario = Basic scenario, so that there is a concrete goal to work towards,
• your agent name, so that it selects the correct agent,
• max_episodes = 1, so that it does not run again and again.

In your challenge directory create a new python file (let's say test.py) with the following contents:

from aica_challenge_1 import Challenge

Challenge().execute(
    "test_run", 
    single_process=True  # No parallel runs and do not supress standard output
)

Note

You could run it from TUI as well, but assuming that you are in the IDE and want to play with your agent, this will make it easier to work with.

Provided that you correctly prepared and saved the run specification, this will start the basic scenario with your agent, which consequently calls its __init__() function.

So... let's start with some introspection into the actions that are available to the agent (still in the __init__()):

for action in self._actions:
  print(f"Action ID: {action.id}\n"
        f"Description: {action.description}\n"
        f"Parameters: {[x.name for x in action.parameters.values()]}\n")

When you execute test.py, you should see something like this:

Action ID: ac1:inspect
Description: Inspecting a machine, where you have an access to.
Parameters: []

Action ID: ac1:scan_host
Description: Scanning of one specific host
Parameters: []

Action ID: ac1:scan_network
Description: Scanning of a subnet
Parameters: ['net']

Action ID: ac1:access_target
Description: Use either exploit or a valid authentication/authorization to open a session to a target service
Parameters: []

Action ID: ac1:exfiltrate_data
Description: Exfiltrate an interesting data from a target.
Parameters: ['path']

If you do, you are sure that:

• your challenge environment is properly set up,
• your agent is correctly registered into the system,
• your run was properly set up,

Therefore, it is now time to truly start building the pseudorandom agent...

First action

We will start by letting the agent feel its closest environment. Check, where it resides and use it as a base for future decisions.

This time, we append the run() function:

async def run(self):
  request = self._msg.create_request(dst_ip="127.0.0.1", 
                                     dst_service="", 
                                     action=self._res.action_store.get("ac1:inspect"))
  self._msg.send_message(request)

We will use the ac1:inspect action that enables you to get information about a machine that you have access to.

You first have to create a request [lines 2-4] that caries the action information. It has three required parameters - dst_ip, dst_service, and action. Considering that you want to inspect a machine your agent is on, then you set dst_ip to a localhost. As there is no service this action is intended for, then you leave dst_service as an empty string. And finally, we use the "ac1:inspect" action. We kept actions internally as list for future, so we do a direct query to the action store once.

With the request ready, we use the send_message() [line 5] function to send it to the system for processing. We then get the response in message_process() function:

async def process_message(self, message: Message) -> Tuple[bool, int]:
    if message.type == MessageType.SIGNAL:
       return
    response = message.cast_to(Response)
    if response.action.id == "ac1:inspect" and response.status.value == StatusValue.SUCCESS and not self._targets:
        for ip in response.content["ips"]:
            for target_ip in IPNetwork(f"{ip}/28").iter_hosts():
                self._targets.append(target_ip)

    return True, 0

Behind the scenes actions

Add self._targets = [] to the __init__ method of your agent.
Add import from cyst.api.environment.message import StatusValue.
Add import from netaddr import IPNetwork.

The process_message() function is the entrypoint for any message directed at the agent. An agent can receive Requests, Responses, Timeouts, Resources, and Signals. All these have a superclass Message.

If you were to create a full-fledged agent, you should be making sure that you correctly process them all. However, Timeouts are a response to agent's explicit call for timeout and Resources to an explicit call for an external resource. None of this is relevant for our scenario. The same goes also for Requests, because an agent can receive a Requests only in the multi-agent scenario. As it is alone in the environment, it will be only processing Responses to its own Requests and an occasional Signal. Therefore, we first guard against Signals (we will discuss them later) on [lines 2-3] and take a shortcut on [line 4] and cast the Message to a Response (which, given Python's type system, is here only to help your IDE with code completion).

On [line 5] we do a specific processing for the ac1:inspect action, as we use it to choose our potential targets. We also want to do this only once and when it is successful (it should be). So we add a guard on message response code and contents of self._targets.

Note

A Response carries the StatusValue.SUCCESS only when the action completed sucessfully and fully. Therefore, you need not to care about StatusOrigin and StatusDetail in such case. However, when the StatusValue is something else, those two can give you a valuable insight on what has happened. More on that in the API reference

According to the ac1:inspect action documentation, a successful response carries two entries in its content:

• ips
• services

As services running on the initial machine are of no importance for our agent, we iterate only over the machine IP addresses [line 6] and decide (a rather arbitrary decision) to use /28 networks [line 7] of these addresses as potential targets [line 8]. This si to make the randomness not explode much into our face.

Now, if you were to run this code, it would store the targets and then just ended. Therefore, it is time to maintain the ...

Action cycle

To make things simple, we will let our agent do only one action at the time. So when it sends a request, it processes the response and then sends another request. And because this is a pseudorandom agent, we will minimize the logic behind its decisions. That will be, after all, your task.

The agent will be operating in the following way:

1. It chooses an action at random,
2. it selects a target,
3. it chooses action parameters,
4. it sends the request.

5. It processes the response and stores new values from it - namely:

   • services,
   • sessions,
   • authorizations.

6. If it finds that it successfully exfiltrated some data, it will cease its activity.

And because we ended the text of the previous section in point 5., we will continue from there.

Processing responses

First of all, services... they are easy. The underlying CYST platform currently does not operate with port numbers, but is instead using string descriptions. This way it sidesteps the issue of using different port numbers for known services. Therefore, the agent cannot choose a random port number, but has to rely on scanning and information gathering. You could create a list of known services and work with that, but the following approach probably has a better chance of succeeding.

Let's continue with the process_message() function...

if response.action.id == "ac1:scan_network":
   for host_report in response.content["success"]:
       self._services.update([x[0] for x in host_report["services"]])
if "services" in response.content:
   self._services.update([x[0] for x in response.content["services"]])

if "data" in response.content:
   self._data.extend(response.content["data"])

self._services = set() was added to the __init__() function behind the scenes. We use it to store all known service names in the infrastructure. As the result is a tuple of a service name and its version, we store only the name. We could have stored the information for each discovered IP, but we would be dangerously crossing the border between pseudorandom and heuristic agent.

As both ac1:scan_host and ac1:scan_network return service names but in different structures, we need to make sure the names are extracted correctly for both actions.

The same goes for 'data' that can be present in the result of the ac1:inspect action and points to paths with known interesting data. self._data = [] was added to the __init__().

Next are the sessions. A session represents an open connection to the target, which enables to either access the target or to use it as a stepping stone for another attack. Terminology-wise, it is an equivalent to Metasploit sessions, and you would generally acquire it through some exploitation or by logging via SSH, for example.

A session gives you several main benefits:

• It enables you to sidestep firewall access controls. If you have an active session to a target, you can communicate with it, regardless of the firewall rules along the way. This is especially useful in scenarios like phishing, or physical token-based infestation, where a reverse shell is opened from the inside, relying on much more lenient outgoing firewall configurations.
• It enables you to hide your source IP address. If you attack through a session, it will appear as if you were attacking from the termination point of the session.
• It enables to you to hop between networks. If you have a session to a machine with multiple IP addresses, each leading to separate network (for example a machine inside DMZ), you can use the session to get access to a different network that wouldn't even be routable otherwise.

If you acquire a session, it makes sense to try and scan the network from its target. You never know which new places it will reveal to you. So, for our agent, we will store each session that terminates at a new place.

if response.session and response.session.end[0] not in self._sessions:
    self._sessions[response.session.end[0]] = response.session

As usual, self._sessions: Dict[IPAddress, Session] = {} was added to the __init__(). On [line 1] we check if the response has any session with it, and if its termination point is new, we save it for later. The .end property of the session is a tuple of the form (IPAddress, service_name), but we do not really need to care about the termination service at this point. Please note that if you do not save the session object, it is as if you closed the session by the end of your action.

And finally authentication and authorization tokens... The underlying CYST platform provides a rather robust authentication and authorization framework that supports a simulation of multifactor authentication and federated authorization. However, in this challenge, we will be using only a limited subset of its capabilities and if you are all too good, we will introduce it later in the challenge and provide an appropriate tutorial.

Currently, we will focus on two use-cases:

• your agent received an authorization through a service exploitation,
• your agent found an authentication token on some machine.

In both cases, you will receive the authentication or authorization token as a part of the response. An authorization token is bound to a service and a target machine, so we will use it accordingly. An authentication token may work wherever an authentication token of that type is accepted, so we will be using it indiscriminately.

if response.auth:
    if isinstance(response.auth, Authorization):
        self._authorizations[(str(response.dst_ip), response.dst_service)] = response.auth
    elif isinstance(response.auth, AuthenticationToken):
        self._auth_tokens.append(response.auth)

The following was added to __init__():

self._authorizations: Dict[Tuple[str, str], Authorization] = {}
self._auth_tokens: List[AuthenticationToken] = []

Here is the code so far, so that you see everything in context:

Agent's code 1/2

import logging

from abc import ABC, abstractmethod
from typing import List, Tuple, Optional, Dict, Any, Union, Callable, Awaitable

from cyst.api.logic.action import Action
from cyst.api.logic.access import Authorization, AuthenticationToken
from cyst.api.environment.environment import EnvironmentMessaging
from cyst.api.environment.message import Request, Response, MessageType, Message, StatusValue
from cyst.api.environment.resources import EnvironmentResources
from cyst.api.network.session import Session
from cyst.api.host.service import ActiveService, ActiveServiceDescription, Service

from netaddr import IPNetwork
from netaddr.ip import IPAddress


class MyAgent(ActiveService):
    def __init__(self, msg: EnvironmentMessaging, res: EnvironmentResources, id: str, args: Optional[Dict[str, Any]]):
        self._msg = msg
        self._res = res
        self._id = id
        self._args = args

        self._actions = res.action_store.get_prefixed("ac1")

        self._targets = []
        self._data = []

        self._services = set()
        self._sessions: Dict[IPAddress, Session] = {}
        self._authorizations: Dict[Tuple[str, str], Authorization] = {}
        self._auth_tokens: List[AuthenticationToken] = []

        for action in self._actions:
            print(f"Action ID: {action.id}\n"
                  f"Description: {action.description}\n"
                  f"Parameters: {[x.name for x in action.parameters.values()]}\n")

    async def run(self):
        request = self._msg.create_request(dst_ip="127.0.0.1",
                                           dst_service="",
                                           action=self._res.action_store.get("ac1:inspect"))
        self._msg.send_message(request)

    async def process_message(self, message: Message) -> Tuple[bool, int]:
        if message.type == MessageType.SIGNAL:
            return
        response = message.cast_to(Response)
        if response.action.id == "ac1:inspect" and response.status.value == StatusValue.SUCCESS and not self._targets:
            for ip in response.content["ips"]:
                for target_ip in IPNetwork(f"{ip}/28").iter_hosts():
                    self._targets.append(target_ip)

        if response.action.id == "ac1:scan_network":
            for host_report in response.content["success"]:
                self._services.update([x[0] for x in host_report["services"]])
        if "services" in response.content:
            self._services.update([x[0] for x in response.content["services"]])

        if "data" in response.content:
            self._data.extend(response.content["data"])

        if response.session and response.session.end[0] not in self._sessions:
            self._sessions[response.session.end[0]] = response.session

        if response.auth:
            if isinstance(response.auth, Authorization):
                self._authorizations[(str(response.dst_ip), response.dst_service)] = response.auth
            elif isinstance(response.auth, AuthenticationToken):
                self._auth_tokens.append(response.auth)

        return True, 0


def create_agent(msg: EnvironmentMessaging, res: EnvironmentResources, id: str,
                 args: Optional[Dict[str, Any]]) -> ActiveService:
    actor = MyAgent(msg, res, id, args)
    return actor


service_description = ActiveServiceDescription(
    "my-awesome-agent",
    "A description of an agent",
    create_agent
)

Preparing requests

Note

In the following sections we are still in the context of the process_message(). However, the code is simplified for the sake of explanation. You should just read it and then copy the example at the end of the page.

We now skip the point 6. (ending the run) and go back to 1. to focus on how to prepare the requests from what the agent has. We actually make it very easy, because randomness...

1. Choosing the action:

action = random.choice(self._actions)

2. Choosing the target:

target_ip = random.choice(self._targets)
target_service = random.choice(tuple(self._services)) if self._services else ""

For the sake of the example, we do not care about the efficiency. So, please forgive the sin of random choice over a set. Also, the target may not have that service at all. But that is one of the risks of being a random agent.

3. Choosing action parameters:

exploits = self._res.exploit_store.get_exploit(service=target_service)
if exploits:
   action.set_exploit(random.choice(exploits))

auth = self._authorizations.get((target_ip, target_service), None)
if not auth and self._auth_tokens:
   auth = random.choice(self._auth_tokens)

session = None
if self._sessions:
    session = random.choice(list(self._sessions.values()))

if "net" in action.parameters:
    action.parameters["net"].value = IPNetwork(f"{target_ip}/29")

if "path" in action.parameters:
    action.parameters["path"].value = random.choice(self._data) if self._data else "."

We check for an exploit for the desired service. [line 1] If it is there, then we use it [line 3]. We do not care much about its parameters. Given enough time, it will work out. If we have an authorization for a specific target service, we use it as well [line 5]. If not, we use an authentication token, if we have one [lines 6-7]. If an action requires a net parameter, we provide it, but choose a very small network around the selected target. Otherwise, the agent could end up scanning like crazy. And in the same vein, if the agent requires the path parameter [lines 16-17], we provide it if we have some data leads and otherwise just use some default.

4. Sending the request:

Now we only have to put all things together and send it.

request = self._msg.create_request(dst_ip=target_ip,
                                   dst_service=target_service,
                                   action=action,
                                   session=session,
                                   auth=auth)

self._msg.send_message(request)

If you were to run it, the agent would be sending requests, processing responses, and it would all run until hitting the action count threshold that you set in the run specification. So, as a final touch, we get back to the point 6. and give the agent a way to report an (un-)successful finish.

Finishing an episode

There are two ways, how the agent can end its activities. Either it gets shut down by the environment or it reaches its goal. As we have mentioned earlier, CYST is a multi-agent platform and there will very likely be multi-agent scenarios in the challenge. In a multi-agent setting, one agent should not be able to decide for the episode to end, just because it is done. The rest of the agents may still be doing something important.

Therefore, CYST adopts a signalling approach to communicate state changes between the environment and the components. Note that this is a signalling strictly between the environment and the components and no signals can be exchanged between components (e.g., agents). That's what messages are for.

So, let's first focus on the latter case, when the agent reaches its goal. For the Basic scenario, the goal is to exfiltrate some data from the target. So, we just need to check for the success of the ac1:exfiltrate_data action.

if response.action.id == "ac1:exfiltrate_data" and response.status.value == StatusValue.SUCCESS:
    self._signal = self._msg.create_signal(signal_origin=self._id, 
                                           state=ComponentState.FINISHED | ComponentState.GOAL_REACHED, 
                                           effect_origin=self._id,
                                           effect_message=response.id)

If that happens, we create a signal that says that the signal is produced by the agent (signal_origin), that it entered a state in which it has finished and the goal was reached (state), that the signal was created on behalf of agent (effect_origin) and that it was the particular response that prompted the state change (effect_message).

You can review the other parameters in the documentation.

We store the signal to use it later and get back to the situation that the environment tells us to finish. In such case there is not much to do. Everything is likely shutting down and our messages will no longer be processed. So the only thing we can do is to signal back that we were not able to reach the goal.

We are going back to the beginning of the message_process() function and for the sake of brevity don't do any complex signal processing, because we abuse the knowledge that in our case, no other signal than a shutdown signal from the environment will come.

if isinstance(message, Signal):
    self._signal = self._msg.create_signal(signal_origin=self._id,
                                           state=ComponentState.TERMINATED | ComponentState.GOAL_UNREACHABLE,
                                           effect_origin=message.effect_origin,
                                           effect_message=message.id)

We are basically saying that the signal is from the agent (signal_origin) who had to enter the terminated state without reaching the goal (state), and it's all environment's fault (effect_origin and effect_message). Know that you do not have to send this signal, as the environment will happily terminate with or without your reaction. But if you want to have this information for further analysis, it's better to put it there.

Now we only need to send the signal and terminate message processing without executing any other action:

if self._signal:
    self._msg.send_message(self._signal)
    return True, 0

And that's it... you have built a pseudorandom agent that can crack in some time the easier scenarios and made yourself acquainted with the basic concepts. And when we say some time, we mean it. For the basic scenario, the reasonable option is to limit it to 1000 virtual seconds and 5000 actions. And even then it occasionally does not finish. You can do better, right?

You can now continue running it from the test script, get back to run configuration to prepare some better runs or test it on more variants, or let it run once and analyze the result.

Here's the code for good measure:

Agent's code 2/2

import random

from typing import List, Tuple, Optional, Dict, Any

from cyst.api.environment.environment import EnvironmentMessaging
from cyst.api.environment.message import Response, Message, StatusValue, Signal, ComponentState
from cyst.api.environment.resources import EnvironmentResources
from cyst.api.host.service import ActiveService, ActiveServiceDescription
from cyst.api.logic.access import Authorization, AuthenticationToken
from cyst.api.network.session import Session

from netaddr import IPNetwork, IPAddress


class MyAgent(ActiveService):
    def __init__(self, msg: EnvironmentMessaging, res: EnvironmentResources, id:str, args: Optional[Dict[str, Any]]):
        self._msg = msg
        self._res = res
        self._id = id
        self._args = args

        self._actions = res.action_store.get_prefixed("ac1")

        self._targets = []
        self._data = []
        self._signal = None

        self._services = set()
        self._sessions: Dict[IPAddress, Session] = {}
        self._authorizations: Dict[Tuple[str, str], Authorization] = {}
        self._auth_tokens: List[AuthenticationToken] = []

        # Print the list of the available actions
        for action in self._actions:
            print(f"Action ID: {action.id}\n"
                  f"Description: {action.description}\n"
                  f"Parameters: {[x.name for x in action.parameters.values()]}\n")

    async def run(self):
        # Inspect the machine the agent is running on
        request = self._msg.create_request(dst_ip="127.0.0.1",
                                           dst_service="",
                                           action=self._res.action_store.get("ac1:inspect"))
        self._msg.send_message(request)

    async def process_message(self, message: Message) -> Tuple[bool, int]:
        # Checking for the termination signal from the environment
        if isinstance(message, Signal):
            self._signal = self._msg.create_signal(signal_origin=self._id,
                                                   state=ComponentState.TERMINATED | ComponentState.GOAL_UNREACHABLE,
                                                   effect_origin=message.effect_origin,
                                                   effect_message=message.id)

        else:
            # No signal, we can cast the message to response
            response = message.cast_to(Response)

            # Checking for goal condition
            if response.action.id == "ac1:exfiltrate_data" and response.status.value == StatusValue.SUCCESS:
                self._signal = self._msg.create_signal(signal_origin=self._id,
                                                       state=ComponentState.FINISHED | ComponentState.GOAL_REACHED,
                                                       effect_origin=self._id,
                                                       effect_message=response.id)

        # We have signal to send. Send it and terminate.
        if self._signal:
            self._msg.send_message(self._signal)
            return True, 0

        # We use inspect action on ourselves to get the list of targets
        if response.action.id == "ac1:inspect" and response.status.value == StatusValue.SUCCESS and not self._targets:
            for ip in response.content["ips"]:
                for target_ip in IPNetwork(f"{ip}/28").iter_hosts():
                    self._targets.append(target_ip)

        # We check for specific data in responses to later populate action parameters
        # scan of a network
        if response.action.id == "ac1:scan_network":
            for host_report in response.content["success"]:
                self._services.update([x[0] for x in host_report["services"]])
        # scan of a host
        if "services" in response.content:
            self._services.update([x[0] for x in response.content["services"]])

        if "data" in response.content:
            self._data.extend(response.content["data"])

        # We store each new session acquired
        if response.session and response.session.end[0] not in self._sessions:
            self._sessions[response.session.end[0]] = response.session

        # And do the same for authorizations and authentications
        if response.auth:
            if isinstance(response.auth, Authorization):
                self._authorizations[(str(response.dst_ip), response.dst_service)] = response.auth
            elif isinstance(response.auth, AuthenticationToken):
                self._auth_tokens.append(response.auth)

        # With the response processing done, create a new request and fill it with random data
        action = random.choice(self._actions)
        target_ip = random.choice(self._targets)
        target_service = random.choice(tuple(self._services)) if self._services else ""

        # Use an exploit if we have some
        exploits = self._res.exploit_store.get_exploit(service=target_service)
        if exploits:
            action.set_exploit(random.choice(exploits))

        # Use an auth if we have some
        auth = self._authorizations.get((target_ip, target_service), None)
        if not auth and self._auth_tokens:
            auth = random.choice(self._auth_tokens)

        # Use a random session
        session = None
        if self._sessions:
            session = random.choice(list(self._sessions.values()))

        # Fill in action parameters if they are required
        if "net" in action.parameters:
            action.parameters["net"].value = IPNetwork(f"{target_ip}/29")

        if "path" in action.parameters:
            action.parameters["path"].value = random.choice(self._data) if self._data else "."

        # And finally craft and send teh request
        request = self._msg.create_request(dst_ip=target_ip,
                                           dst_service=target_service,
                                           action=action,
                                           session=session,
                                           auth=auth)

        self._msg.send_message(request)

        return True, 0


def create_agent(msg: EnvironmentMessaging, res: EnvironmentResources, id:str, args: Optional[Dict[str, Any]]) -> ActiveService:
    actor = MyAgent(msg, res, id, args)
    return actor


service_description = ActiveServiceDescription(
    "my-awesome-agent",
    "A description of an agent",
    create_agent
)