Execute data manipulation bots from agent
This commit is contained in:
@@ -24,6 +24,20 @@ class AgentExecutionDefinition(BaseModel):
|
|||||||
data_manipulation_p_of_success: float = 0.1
|
data_manipulation_p_of_success: float = 0.1
|
||||||
"The probability of data manipulation succeeding."
|
"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):
|
class AgentStartSettings(BaseModel):
|
||||||
"""Configuration values for when an agent starts performing actions."""
|
"""Configuration values for when an agent starts performing actions."""
|
||||||
@@ -42,6 +56,20 @@ class AgentSettings(BaseModel):
|
|||||||
start_settings: Optional[AgentStartSettings] = None
|
start_settings: Optional[AgentStartSettings] = None
|
||||||
"Configuration for when an agent begins performing it's actions"
|
"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):
|
class AbstractAgent(ABC):
|
||||||
"""Base class for scripted and RL agents."""
|
"""Base class for scripted and RL agents."""
|
||||||
@@ -149,6 +177,8 @@ class RandomAgent(AbstractScriptedAgent):
|
|||||||
class DataManipulationAgent(AbstractScriptedAgent):
|
class DataManipulationAgent(AbstractScriptedAgent):
|
||||||
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
|
"""Agent that uses a DataManipulationBot to perform an SQL injection attack."""
|
||||||
|
|
||||||
|
data_manipulation_bots: List["DataManipulationBot"] = []
|
||||||
|
|
||||||
def __init__(self, *args, **kwargs):
|
def __init__(self, *args, **kwargs):
|
||||||
super().__init__(*args, **kwargs)
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
@@ -163,6 +193,7 @@ class DataManipulationAgent(AbstractScriptedAgent):
|
|||||||
|
|
||||||
if bot_sw is not None:
|
if bot_sw is not None:
|
||||||
bot_sw.execution_definition = self.execution_definition
|
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]:
|
def get_action(self, obs: ObsType, reward: float = None) -> Tuple[str, Dict]:
|
||||||
"""Randomly sample an action from the action space.
|
"""Randomly sample an action from the action space.
|
||||||
@@ -174,7 +205,12 @@ class DataManipulationAgent(AbstractScriptedAgent):
|
|||||||
:return: _description_
|
:return: _description_
|
||||||
:rtype: Tuple[str, Dict]
|
: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):
|
class AbstractGATEAgent(AbstractAgent):
|
||||||
|
|||||||
@@ -444,8 +444,8 @@ class PrimaiteSession:
|
|||||||
# CREATE REWARD FUNCTION
|
# CREATE REWARD FUNCTION
|
||||||
rew_function = RewardFunction.from_config(reward_function_cfg, session=sess)
|
rew_function = RewardFunction.from_config(reward_function_cfg, session=sess)
|
||||||
|
|
||||||
execution_definition = AgentExecutionDefinition(**agent_cfg.get("execution_definition", {}))
|
execution_definition = AgentExecutionDefinition.from_config(agent_cfg.get("execution_definition"))
|
||||||
agent_settings = AgentSettings(**agent_cfg.get("agent_settings", {}))
|
agent_settings = AgentSettings.from_config(agent_cfg.get("agent_settings"))
|
||||||
|
|
||||||
# CREATE AGENT
|
# CREATE AGENT
|
||||||
if agent_type == "GreenWebBrowsingAgent":
|
if agent_type == "GreenWebBrowsingAgent":
|
||||||
|
|||||||
@@ -130,7 +130,7 @@ class DatabaseClient(Application):
|
|||||||
|
|
||||||
def execute(self) -> None:
|
def execute(self) -> None:
|
||||||
"""Run the DatabaseClient."""
|
"""Run the DatabaseClient."""
|
||||||
super().execute()
|
# super().execute()
|
||||||
if self.operating_state == ApplicationOperatingState.RUNNING:
|
if self.operating_state == ApplicationOperatingState.RUNNING:
|
||||||
self.connect()
|
self.connect()
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user