def run_agent(self, workspace_name, agent_id):
"""
---
tags: ["Agent"]
description: Runs an agent
responses:
400:
description: Bad request
201:
description: Ok
content:
application/json:
schema: AgentSchema
"""
if flask.request.content_type != 'application/json':
abort(400, "Only application/json is a valid content-type")
data = self._parse_data(AgentRunSchema(unknown=EXCLUDE), request)
agent = self._get_object(agent_id, workspace_name)
workspace = self._get_workspace(workspace_name)
executor_data = data['executorData']
try:
executor = Executor.query.filter(
Executor.name == executor_data['executor'],
Executor.agent_id == agent_id).one()
agent_execution = AgentExecution(
running=None,
successful=None,
message='',
executor=executor,
workspace_id=workspace.id,
parameters_data=executor_data["args"])
db.session.add(agent_execution)
db.session.commit()
changes_queue.put({
'execution_id': agent_execution.id,
'agent_id': agent.id,
'workspace': agent_execution.workspace.name,
'action': 'RUN',
"executor": executor_data.get('executor'),
"args": executor_data.get('args')
})
except NoResultFound as e:
logger.exception(e)
abort(400, "Can not find an agent execution with that id")
return flask.jsonify({
'successful': True,
})
def run_agent(self, workspace_name, agent_id):
"""
---
tags: ["Agent"]
description: Runs an agent
responses:
400:
description: Bad request
201:
description: Ok
content:
application/json:
schema: AgentSchema
"""
if flask.request.content_type != 'application/json':
abort(400, "Only application/json is a valid content-type")
data = self._parse_data(AgentRunSchema(unknown=EXCLUDE), request)
agent = self._get_object(agent_id, workspace_name)
workspace = self._get_workspace(workspace_name)
executor_data = data['executorData']
try:
executor = Executor.query.filter(Executor.name == executor_data['executor'],
Executor.agent_id == agent_id).one()
# VALIDATE
errors = {}
for param_name, param_data in executor_data["args"].items():
if executor.parameters_metadata.get(param_name):
val_error = type_validate(executor.parameters_metadata[param_name]['type'], param_data)
if val_error:
errors[param_name] = val_error
else:
errors['message'] = f'"{param_name}" not recognized as an executor argument'
if errors:
response = jsonify(errors)
response.status_code = 400
abort(response)
params = ', '.join([f'{key}={value}' for (key, value) in executor_data["args"].items()])
command = Command(
import_source="agent",
tool=agent.name,
command=executor.name,
user='',
hostname='',
params=params,
start_date=datetime.utcnow(),
workspace=workspace
)
agent_execution = AgentExecution(
running=None,
successful=None,
message='',
executor=executor,
workspace_id=workspace.id,
parameters_data=executor_data["args"],
command=command
)
executor.last_run = datetime.utcnow()
db.session.add(agent_execution)
db.session.commit()
changes_queue.put({
'execution_id': agent_execution.id,
'agent_id': agent.id,
'workspace': agent_execution.workspace.name,
'action': 'RUN',
"executor": executor_data.get('executor'),
"args": executor_data.get('args')
})
except NoResultFound as e:
logger.exception(e)
abort(400, "Can not find an agent execution with that id")
else:
return flask.jsonify({
'command_id': command.id,
})
def onMessage(self, payload, is_binary):
from faraday.server.web import app # pylint:disable=import-outside-toplevel
"""
We only support JOIN and LEAVE workspace messages.
When authentication is implemented we need to verify
that the user can join the selected workspace.
When authentication is implemented we need to reply
the client if the join failed.
"""
if not is_binary:
message = json.loads(payload)
if message['action'] == 'JOIN_WORKSPACE':
if 'workspace' not in message or 'token' not in message:
logger.warning('Invalid join workspace message: '
'{}'.format(message))
self.sendClose()
return
signer = itsdangerous.TimestampSigner(app.config['SECRET_KEY'],
salt="websocket")
try:
workspace_id = signer.unsign(message['token'], max_age=60)
except itsdangerous.BadData as e:
self.sendClose()
logger.warning('Invalid websocket token for workspace '
'{}'.format(message['workspace']))
logger.exception(e)
else:
with app.app_context():
workspace = Workspace.query.get(int(workspace_id))
if workspace.name != message['workspace']:
logger.warning(
'Trying to join workspace {} with token of '
'workspace {}. Rejecting.'.format(
message['workspace'], workspace.name))
self.sendClose()
else:
self.factory.join_workspace(self, message['workspace'])
if message['action'] == 'LEAVE_WORKSPACE':
self.factory.leave_workspace(self, message['workspace'])
if message['action'] == 'JOIN_AGENT':
if 'token' not in message or 'executors' not in message:
logger.warn("Invalid agent join message")
self.state = WebSocketProtocol.STATE_CLOSING
self.sendClose()
return False
with app.app_context():
try:
agent = decode_agent_websocket_token(message['token'])
update_executors(agent, message['executors'])
except ValueError:
logger.warn('Invalid agent token!')
self.state = WebSocketProtocol.STATE_CLOSING
self.sendClose()
return False
# factory will now send broadcast messages to the agent
return self.factory.join_agent(self, agent)
if message['action'] == 'LEAVE_AGENT':
with app.app_context():
(agent_id, ) = [
k for (k, v) in connected_agents.items() if v == self
]
agent = Agent.query.get(agent_id)
assert agent is not None # TODO the agent could be deleted here
self.factory.leave_agent(self, agent)
self.state = WebSocketProtocol.STATE_CLOSING
self.sendClose()
return False
if message['action'] == 'RUN_STATUS':
with app.app_context():
if 'executor_name' not in message:
logger.warning(
'Missing executor_name param in message: '
'{}'.format(message))
self.sendClose()
return
(agent_id, ) = [
k for (k, v) in connected_agents.items() if v == self
]
agent = Agent.query.get(agent_id)
assert agent is not None # TODO the agent could be deleted here
executor = Executor.query.filter(
Executor.name == message['executor_name'],
Executor.agent_id == agent_id).first()
if executor:
successful = message.get('successful', None)
running = message.get('running', None)
msg = message['message']
agent_execution = AgentExecution(
running=running,
successful=successful,
message=msg,
executor=executor,
workspace_id=executor.agent.workspace_id)
db.session.add(agent_execution)
db.session.commit()