class Executor:
__control_dict = {
Sections.EXECUTOR_DATA: {
"cmd": control_str(True),
"repo_executor": control_str(True),
"max_size": control_int(True),
}
}
def __init__(self, name: str, config):
name = name.strip()
self.control_config(name, config)
self.name = name
executor_section = Sections.EXECUTOR_DATA.format(name)
params_section = Sections.EXECUTOR_PARAMS.format(name)
varenvs_section = Sections.EXECUTOR_VARENVS.format(name)
self.repo_name = config[executor_section].get("repo_executor", None)
if self.repo_name:
metadata = executor_metadata(self.repo_name)
repo_path = executor_folder() / self.repo_name
self.cmd = metadata["cmd"].format(EXECUTOR_FILE_PATH=repo_path)
else:
self.cmd = config[executor_section].get("cmd")
self.max_size = int(config[executor_section].get("max_size", 64 * 1024))
self.params = dict(config[params_section]) if params_section in config else {}
self.params = {key: value.lower() in ["t", "true"] for key, value in self.params.items()}
self.varenvs = dict(config[varenvs_section]) if varenvs_section in config else {}
def control_config(self, name, config):
if " " in name:
raise ValueError("Executor names can't contains space character, passed name:" f"{name}")
if Sections.EXECUTOR_DATA.format(name) not in config:
raise ValueError(f"{name} is an executor name but there is no proper section")
for section in self.__control_dict:
for option in self.__control_dict[section]:
value = config.get(section.format(name), option) if option in config[section.format(name)] else None
self.__control_dict[section][option](option, value)
params_section = Sections.EXECUTOR_PARAMS.format(name)
if params_section in config:
for option in config[params_section]:
value = config.get(params_section, option)
control_bool(option, value)
async def check_cmds(self):
if self.repo_name is None:
return True
metadata = executor_metadata(self.repo_name)
if not await check_commands(metadata):
logger.info(
f"{Bcolors.WARNING}Invalid bash dependency for " f"{Bcolors.BOLD}{self.repo_name}{Bcolors.ENDC}"
)
return False
else:
return True
class Executor:
__control_dict = {
"cmd": control_str(True),
"repo_executor": control_str(True),
"max_size": control_int(True),
}
def __init__(self, name: str, config):
name = name.strip()
self.control_config(name, config)
self.name = name
self.repo_executor = config.get("repo_executor")
if self.repo_executor:
self.repo_name = re.search(r"(^[a-zA-Z0-9_-]+)(?:\..*)*$",
self.repo_executor).group(1)
metadata = executor_metadata(self.repo_name)
repo_path = executor_folder() / self.repo_executor
self.cmd = metadata["cmd"].format(EXECUTOR_FILE_PATH=repo_path)
else:
self.cmd = config.get("cmd")
self.max_size = int(config.get("max_size", 64 * 1024))
self.params = dict(config[Sections.EXECUTOR_PARAMS]
) if Sections.EXECUTOR_PARAMS in config else {}
self.varenvs = dict(config[Sections.EXECUTOR_VARENVS]
) if Sections.EXECUTOR_VARENVS in config else {}
def control_config(self, name, config):
if " " in name:
raise ValueError(
"Executor names can't contains space character, passed name:"
f"{name}")
for option in self.__control_dict:
value = config[option] if option in config else None
self.__control_dict[option](option, value)
if Sections.EXECUTOR_PARAMS in config:
value = config.get(Sections.EXECUTOR_PARAMS)
errors = ParamsSchema().validate({"params": value})
if errors:
raise ValueError(errors)
async def check_cmds(self):
if self.repo_executor is None:
return True
repo_name = re.search(r"(^[a-zA-Z0-9_-]+)(?:\..*)*$",
self.repo_executor).group(1)
metadata = executor_metadata(repo_name)
if not await check_commands(metadata):
logger.info(f"{Bcolors.WARNING}Invalid bash dependency for "
f"{Bcolors.BOLD}{self.repo_name}{Bcolors.ENDC}")
return False
else:
return True
class Executor:
__control_dict = {
Sections.EXECUTOR_DATA: {
"cmd": control_str(),
"max_size": control_int(True)
}
}
def __init__(self, name: str, config):
name = name.strip()
self.control_config(name, config)
self.name = name
executor_section = Sections.EXECUTOR_DATA.format(name)
params_section = Sections.EXECUTOR_PARAMS.format(name)
varenvs_section = Sections.EXECUTOR_VARENVS.format(name)
self.cmd = config.get(executor_section, "cmd")
self.max_size = int(config[executor_section].get(
"max_size", 64 * 1024))
self.params = dict(
config[params_section]) if params_section in config else {}
self.params = {
key: value.lower() in ["t", "true"]
for key, value in self.params.items()
}
self.varenvs = dict(
config[varenvs_section]) if varenvs_section in config else {}
def control_config(self, name, config):
if " " in name:
raise ValueError(
f"Executor names can't contains space character, passed name: {name}"
)
if Sections.EXECUTOR_DATA.format(name) not in config:
raise ValueError(
f"{name} is an executor name but there is no proper section")
for section in self.__control_dict:
for option in self.__control_dict[section]:
value = config.get(
section.format(name),
option) if option in config[section.format(name)] else None
self.__control_dict[section][option](option, value)
params_section = Sections.EXECUTOR_PARAMS.format(name)
if params_section in config:
for option in config[params_section]:
value = config.get(params_section, option)
control_bool(option, value)
class Sections:
TOKENS = "tokens"
SERVER = "server"
AGENT = "agent"
EXECUTOR_VARENVS = "{}_varenvs"
EXECUTOR_PARAMS = "{}_params"
EXECUTOR_DATA = "{}"
__control_dict = {
Sections.SERVER: {
"host": control_host,
"ssl": control_bool,
"ssl_cert": control_str(nullable=True),
"api_port": control_int(),
"websocket_port": control_int(),
"workspaces": control_list(can_repeat=False),
},
Sections.TOKENS: {
"registration": control_registration_token,
"agent": control_agent_token,
},
Sections.AGENT: {
"agent_name": control_str(),
"executors": control_list(can_repeat=False),
},
}
def control_config():
class Dispatcher:
__control_dict = {
Sections.SERVER: {
"host": control_host,
"api_port": control_int(),
"websocket_port": control_int(),
"workspace": control_str
},
Sections.TOKENS: {
"registration": control_registration_token,
"agent": control_agent_token
},
Sections.AGENT: {
"agent_name": control_str,
"executors": control_list(can_repeat=False)
},
}
def __init__(self, session, config_path=None):
reset_config(filepath=config_path)
self.control_config()
self.config_path = config_path
self.host = config.get(Sections.SERVER, "host")
self.api_port = config.get(Sections.SERVER, "api_port")
self.websocket_port = config.get(Sections.SERVER, "websocket_port")
self.workspace = config.get(Sections.SERVER, "workspace")
self.agent_token = config[Sections.TOKENS].get("agent", None)
self.agent_name = config.get(Sections.AGENT, "agent_name")
self.session = session
self.websocket = None
self.websocket_token = None
self.executors = {
executor_name: Executor(executor_name, config)
for executor_name in config[Sections.AGENT].get("executors",
[]).split(",")
}
async def reset_websocket_token(self):
# I'm built so I ask for websocket token
headers = {"Authorization": f"Agent {self.agent_token}"}
websocket_token_response = await self.session.post(api_url(
self.host,
self.api_port,
postfix='/_api/v2/agent_websocket_token/'),
headers=headers)
websocket_token_json = await websocket_token_response.json()
return websocket_token_json["token"]
async def register(self):
if self.agent_token is None:
registration_token = self.agent_token = config.get(
Sections.TOKENS, "registration")
assert registration_token is not None, "The registration token is mandatory"
token_registration_url = api_url(
self.host,
self.api_port,
postfix=f"/_api/v2/ws/{self.workspace}/agent_registration/")
logger.info(f"token_registration_url: {token_registration_url}")
try:
token_response = await self.session.post(
token_registration_url,
json={
'token': registration_token,
'name': self.agent_name
})
assert token_response.status == 201
token = await token_response.json()
self.agent_token = token["token"]
config.set(Sections.TOKENS, "agent", self.agent_token)
save_config(self.config_path)
except ClientResponseError as e:
if e.status == 404:
logger.info(
f'404 HTTP ERROR received: Workspace "{self.workspace}" not found'
)
return
else:
logger.info(f"Unexpected error: {e}")
raise e
self.websocket_token = await self.reset_websocket_token()
async def connect(self, out_func=None):
if not self.websocket_token and not out_func:
return
connected_data = json.dumps({
'action':
'JOIN_AGENT',
'workspace':
self.workspace,
'token':
self.websocket_token,
'executors': [{
"executor_name": executor.name,
"args": executor.params
} for executor in self.executors.values()]
})
if out_func is None:
async with websockets.connect(
websocket_url(self.host,
self.websocket_port)) as websocket:
await websocket.send(connected_data)
logger.info("Connection to Faraday server succeeded")
self.websocket = websocket
await self.run_await(
) # This line can we called from outside (in main)
else:
await out_func(connected_data)
async def run_await(self):
while True:
# Next line must be uncommented, when faraday (and dispatcher) maintains the keep alive
data = await self.websocket.recv()
asyncio.create_task(self.run_once(data))
async def run_once(self, data: str = None, out_func=None):
out_func = out_func if out_func is not None else self.websocket.send
logger.info('Parsing data: %s', data)
data_dict = json.loads(data)
if "action" not in data_dict:
logger.info("Data not contains action to do")
await out_func(
json.dumps({
"error":
"'action' key is mandatory in this websocket connection"
}))
return
if data_dict["action"] not in ["RUN"
]: # ONLY SUPPORTED COMMAND FOR NOW
logger.info("Unrecognized action")
await out_func(
json.dumps({
f"{data_dict['action']}_RESPONSE":
"Error: Unrecognized action"
}))
return
if data_dict["action"] == "RUN":
if "executor" not in data_dict:
logger.error("No executor selected")
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"running":
False,
"message":
f"No executor selected to {self.agent_name} agent"
}))
return
if data_dict["executor"] not in self.executors:
logger.error("The selected executor not exists")
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"executor_name":
data_dict['executor'],
"running":
False,
"message":
f"The selected executor {data_dict['executor']} not exists in {self.agent_name} "
f"agent"
}))
return
executor = self.executors[data_dict["executor"]]
params = list(executor.params.keys()).copy()
passed_params = data_dict['args'] if 'args' in data_dict else {}
[params.remove(param) for param in config.defaults()]
all_accepted = all([
any([
param in passed_param # Control any available param
for param in params # was passed
]) for passed_param in passed_params # For all passed params
])
if not all_accepted:
logger.error(
"Unexpected argument passed to {} executor".format(
executor.name))
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"executor_name":
executor.name,
"running":
False,
"message":
f"Unexpected argument(s) passed to {executor.name} executor from {self.agent_name} "
f"agent"
}))
mandatory_full = all([
not executor.params[param] # All params is not mandatory
or any([
param in passed_param
for passed_param in passed_params # Or was passed
]) for param in params
])
if not mandatory_full:
logger.error(
"Mandatory argument not passed to {} executor".format(
executor.name))
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"executor_name":
executor.name,
"running":
False,
"message":
f"Mandatory argument(s) not passed to {executor.name} executor from "
f"{self.agent_name} agent"
}))
if mandatory_full and all_accepted:
running_msg = f"Running {executor.name} executor from {self.agent_name} agent"
logger.info("Running {} executor".format(executor.name))
process = await self.create_process(executor, passed_params)
tasks = [
StdOutLineProcessor(process, self.session).process_f(),
StdErrLineProcessor(process).process_f(),
]
await out_func(
json.dumps({
"action": "RUN_STATUS",
"executor_name": executor.name,
"running": True,
"message": running_msg
}))
await asyncio.gather(*tasks)
await process.communicate()
assert process.returncode is not None
if process.returncode == 0:
logger.info("Executor {} finished successfully".format(
executor.name))
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"executor_name":
executor.name,
"successful":
True,
"message":
f"Executor {executor.name} from {self.agent_name} finished successfully"
}))
else:
logger.warning(
f"Executor {executor.name} finished with exit code {process.returncode}"
)
await out_func(
json.dumps({
"action":
"RUN_STATUS",
"executor_name":
executor.name,
"successful":
False,
"message":
f"Executor {executor.name} from {self.agent_name} failed"
}))
async def create_process(self, executor: Executor, args):
env = os.environ.copy()
if isinstance(args, dict):
for k in args:
env[f"EXECUTOR_CONFIG_{k.upper()}"] = str(args[k])
else:
logger.error("Args from data received has a not supported type")
raise ValueError(
"Args from data received has a not supported type")
for varenv, value in executor.varenvs.items():
env[f"{varenv.upper()}"] = value
process = await asyncio.create_subprocess_shell(
executor.cmd,
stdout=asyncio.subprocess.PIPE,
stderr=asyncio.subprocess.PIPE,
env=env,
limit=executor.max_size
# If the config is not set, use async.io default
)
return process
def control_config(self):
for section in self.__control_dict:
for option in self.__control_dict[section]:
if section not in config:
err = f"Section {section} is an mandatory section in the config" # TODO "run config cmd"
logger.error(err)
raise ValueError(err)
value = config.get(
section, option) if option in config[section] else None
self.__control_dict[section][option](option, value)