Execute data manipulation bots from agent

This commit is contained in:
Jake Walker
2023-11-21 11:42:01 +00:00
parent 7e0e8a4768
commit 2975aa8827
3 changed files with 40 additions and 4 deletions

View File

@@ -24,6 +24,20 @@ class AgentExecutionDefinition(BaseModel):
data_manipulation_p_of_success: float = 0.1
"The probability of data manipulation succeeding."
@classmethod
def from_config(cls, config: Optional[Dict]) -> "AgentExecutionDefinition":
"""Construct an AgentExecutionDefinition from a config dictionary.
:param config: A dict of options for the execution definition.
:type config: Dict
:return: The execution definition.
:rtype: AgentExecutionDefinition
"""
if config is None:
return cls()
return cls(**config)
class AgentStartSettings(BaseModel):
"""Configuration values for when an agent starts performing actions."""
@@ -42,6 +56,20 @@ class AgentSettings(BaseModel):
start_settings: Optional[AgentStartSettings] = None
"Configuration for when an agent begins performing it's actions"
@classmethod
def from_config(cls, config: Optional[Dict]) -> "AgentSettings":
"""Construct agent settings from a config dictionary.
:param config: A dict of options for the agent settings.
:type config: Dict
:return: The agent settings.
:rtype: AgentSettings
"""
if config is None:
return cls()
return cls(**config)
class AbstractAgent(ABC):
"""Base class for scripted and RL agents."""
@@ -149,6 +177,8 @@ class RandomAgent(AbstractScriptedAgent):
class DataManipulationAgent(AbstractScriptedAgent):
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
data_manipulation_bots: List["DataManipulationBot"] = []
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@@ -163,6 +193,7 @@ class DataManipulationAgent(AbstractScriptedAgent):
if bot_sw is not None:
bot_sw.execution_definition = self.execution_definition
self.data_manipulation_bots.append(bot_sw)
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
"""Randomly sample an action from the action space.
@@ -174,7 +205,12 @@ class DataManipulationAgent(AbstractScriptedAgent):
:return: _description_
:rtype: Tuple[str, Dict]
"""
return self.action_space.get_action(self.action_space.space.sample())
# TODO: Move this to the appropriate place
# return self.action_space.get_action(self.action_space.space.sample())
for bot in self.data_manipulation_bots:
bot.execute()
return ("DONOTHING", {"dummy": 0})
class AbstractGATEAgent(AbstractAgent):

View File

@@ -444,8 +444,8 @@ class PrimaiteSession:
# CREATE REWARD FUNCTION
rew_function = RewardFunction.from_config(reward_function_cfg, session=sess)
execution_definition = AgentExecutionDefinition(**agent_cfg.get("execution_definition", {}))
agent_settings = AgentSettings(**agent_cfg.get("agent_settings", {}))
execution_definition = AgentExecutionDefinition.from_config(agent_cfg.get("execution_definition"))
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
# CREATE AGENT
if agent_type == "GreenWebBrowsingAgent":

View File

@@ -130,7 +130,7 @@ class DatabaseClient(Application):
def execute(self) -> None:
"""Run the DatabaseClient."""
super().execute()
# super().execute()
if self.operating_state == ApplicationOperatingState.RUNNING:
self.connect()