def determine_server_port(self, args: Dict) -> int:
"""
Return port to be used by the server.
Will raise a NotFoundException if port is already in use.
When port given in command line, only that port will be tried.
When no port specified will try self.MAX_PORTS_TRIED times,
starting with self.PORT_DEFAULT.
"""
command_line_port = args.get("port")
if command_line_port is not None:
if self.is_port_taken(command_line_port):
raise NotFoundException(
f"Port {command_line_port} already in use, exiting.", )
else:
return command_line_port
ports = [self.PORT_DEFAULT] + randint(
1025,
65536,
self.MAX_PORTS_TRIED - 1,
).tolist()
for port in ports:
if not self.is_port_taken(port):
return port
raise NotFoundException(
f"Unable to find a free port in {len(ports)} attempts, exiting.", )
def verify_file_path(path: str) -> None:
"""Check if path can be accessed."""
restricted_paths = [
"/bin",
"/boot",
"/dev",
"/etc",
"/lib",
"/media",
"/proc",
"/root",
"/run",
"/sbin",
"/snap",
"/srv",
"/swapfile",
"/usr",
"/var",
]
real_path = os.path.realpath(path)
if not os.path.exists(real_path):
raise NotFoundException("File not found.")
if os.stat(real_path).st_uid == 0:
raise AccessDeniedException("Access denied.")
for path_element in real_path.split(os.sep):
if path_element.startswith("."):
raise AccessDeniedException("Access denied.")
for restricted_path in restricted_paths:
if real_path.startswith(restricted_path):
raise AccessDeniedException("Access denied.")
def get_model(self, path: str) -> Model:
"""Get Model for given path."""
for model_type in self.model_types:
if model_type.supports_path(path):
return model_type(path)
raise NotFoundException(f"Unable to find Model for path {path}")
def test_get_output_with_failure(
self,
mocked_get_workload_data: MagicMock,
mocked_serve_from_filesystem: MagicMock,
) -> None:
"""Test get_output."""
mocked_get_workload_data.return_value = {
"log_path": "/some/fake/output/path.log",
}
mocked_serve_from_filesystem.side_effect = NotFoundException(
"Unable to find file.")
actual = WorkloadService.get_output(
{
"workload_id": ["requested workload id"],
}, )
self.assertEqual("Unable to find file.", actual.data.decode("utf-8"))
self.assertEqual(404, actual.status_code)
self.assertEqual("3", actual.headers.get("refresh"))
mocked_get_workload_data.assert_called_with("requested workload id")
mocked_serve_from_filesystem.assert_called_once_with(
path="/some/fake/output/path.log",
mimetype="text/plain",
)
def test_from_exception_for_not_found_exception(self) -> None:
"""Test from_exception for NotFoundException."""
message = "There's nothing here!"
response = ResponseGenerator.from_exception(NotFoundException(message))
self.assertEqual(404, response.status_code)
self.assertEqual(message, response.data.decode("utf-8"))
def test_serve_from_filesystem_passes_verify_file_path_errors(
self,
mocked_verify_file_path: MagicMock,
) -> None:
"""Test that exceptions from verify_file_path are propagated."""
mocked_verify_file_path.side_effect = NotFoundException(
"Expected Error Message")
with self.assertRaisesRegex(NotFoundException,
"Expected Error Message"):
ResponseGenerator.serve_from_filesystem("foo.txt")
def get_code_template(data: dict) -> Response:
"""Get code template file for requested Workload."""
workload_id = RequestDataProcessor.get_string_value(data, "workload_id")
workdir = Workdir()
workload_data = workdir.get_workload_data(workload_id)
code_template_path = workload_data.get("code_template_path")
if not code_template_path:
raise NotFoundException(f"Unable to find code template file for {workload_id}")
return ResponseGenerator.serve_from_filesystem(
path=code_template_path,
mimetype="text/x-python",
)
def delete_workload(self, workload_id: str) -> None:
"""Delete workload by ID."""
if workload_id in self.workloads_data.get("workloads", {}):
try:
workload = self.workloads_data.get("workloads", {}).get(workload_id, {})
workload_path = workload.get("workload_path", None)
del self.workloads_data.get("workloads", {})[workload_id]
self.dump()
if workload_path:
shutil.rmtree(workload_path)
except Exception as e:
raise InternalException(
f"Error while removing workload {workload_id}:\n{str(e)}",
)
else:
raise NotFoundException(f"Can't find workload ID {workload_id}.")
def get_directory_entries(data: Dict[str, Any], ) -> Dict[str, Any]:
"""Get directory entries."""
try:
path = get_requested_path(data)
verify_file_path(path)
contents = get_non_hidden_directory_entries(path)
contents = filter_requested_entries(contents, get_filter_value(data))
return {
"path": path,
"contents": sort_entries(contents),
}
except PermissionError as err:
raise AccessDeniedException(err)
except FileNotFoundError as err:
raise NotFoundException(err)
except NotADirectoryError as err:
raise ClientErrorException(err)
def get_output(data: dict) -> Response:
"""Get config file for requested Workload."""
workload_id = RequestDataProcessor.get_string_value(data, "workload_id")
workdir = Workdir()
workload_data = workdir.get_workload_data(workload_id)
log_path = workload_data.get("log_path")
if not log_path:
raise NotFoundException(f"Unable to find output log for {workload_id}")
try:
response = ResponseGenerator.serve_from_filesystem(
path=log_path,
mimetype="text/plain",
)
except Exception as e:
response = ResponseGenerator.from_exception(e)
return ResponseGenerator.add_refresh(
response=response,
refresh_time=3,
)
def ensure_model_readable(self, model_path: str) -> None:
"""Throw Exception if can't read model."""
if not os.path.isfile(model_path):
raise NotFoundException(f"Unable to find {model_path} model")
def _get_node(self, id: str) -> Node:
"""Get a node by id."""
if id not in self._nodes:
raise NotFoundException(f"Node id: {id} not found in Graph")
return self._nodes[id]