diff --git a/agent/agent.py b/agent/agent.py index 496841b7fb3..158a138005a 100644 --- a/agent/agent.py +++ b/agent/agent.py @@ -4,6 +4,7 @@ import argparse import cgi +import enum import http.server import ipaddress import json @@ -37,7 +38,7 @@ if sys.maxsize > 2**32 and sys.platform == "win32": sys.exit("You should install python3 x86! not x64") -AGENT_VERSION = "0.13" +AGENT_VERSION = "0.14" AGENT_FEATURES = [ "execpy", "execute", @@ -47,13 +48,35 @@ "unicodepath", ] -STATUS_INIT = 0x0001 -STATUS_RUNNING = 0x0002 -STATUS_COMPLETED = 0x0003 -STATUS_FAILED = 0x0004 + +class Status(enum.IntEnum): + INIT = 1 + RUNNING = 2 + COMPLETE = 3 + FAILED = 4 + EXCEPTION = 5 + + def __str__(self): + return f"{self.name.lower()}" + + @classmethod + def _missing_(cls, value): + if not isinstance(value, str): + return None + value = value.lower() + for member in cls: + if str(member) == value: + return member + if value.isnumeric() and int(value) == member.value: + return member + return None + ANALYZER_FOLDER = "" -state = {"status": STATUS_INIT} +state = { + "status": Status.INIT, + "description": "", +} class MiniHTTPRequestHandler(http.server.SimpleHTTPRequestHandler): @@ -197,11 +220,11 @@ def __init__(self, path): self.status_code = 200 def init(self): - if not os.path.isfile(self.path): + if os.path.isfile(self.path) and os.access(self.path, os.R_OK): + self.length = os.path.getsize(self.path) + else: self.status_code = 404 self.length = 0 - else: - self.length = os.path.getsize(self.path) def write(self, sock): if not self.length: @@ -270,15 +293,17 @@ def get_index(): @app.route("/status") def get_status(): - return json_success("Analysis status", status=state.get("status"), description=state.get("description")) + return json_success("Analysis status", status=str(state.get("status")), description=state.get("description")) @app.route("/status", methods=["POST"]) def put_status(): - if "status" not in request.form: - return json_error(400, "No status has been provided") + try: + status = Status(request.form.get("status")) + except ValueError: + return json_error(400, "No valid status has been provided") - state["status"] = request.form["status"] + state["status"] = status state["description"] = request.form.get("description") return json_success("Analysis status updated") @@ -449,11 +474,12 @@ def do_execute(): p = subprocess.Popen(request.form["command"], shell=shell, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr = p.communicate() except Exception: - state["status"] = STATUS_FAILED + state["status"] = Status.FAILED state["description"] = "Error execute command" return json_exception("Error executing command") - state["status"] = STATUS_RUNNING + state["status"] = Status.RUNNING + state["description"] = "" return json_success("Successfully executed command", stdout=stdout, stderr=stderr) @@ -480,11 +506,11 @@ def do_execpy(): p = subprocess.Popen(args, cwd=cwd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) stdout, stderr = p.communicate() except Exception: - state["status"] = STATUS_FAILED + state["status"] = Status.FAILED state["description"] = "Error executing command" return json_exception("Error executing command") - state["status"] = STATUS_RUNNING + state["status"] = Status.RUNNING return json_success("Successfully executed command", stdout=stdout, stderr=stderr) diff --git a/agent/test_agent.py b/agent/test_agent.py index bc38c28f654..cff1711eeab 100644 --- a/agent/test_agent.py +++ b/agent/test_agent.py @@ -35,7 +35,7 @@ class TestAgent: agent_process: multiprocessing.Process = None def setup_method(self): - agent.state = {"status": agent.STATUS_INIT, "description": "", "async_subprocess": None} + agent.state = {"status": agent.Status.INIT, "description": "", "async_subprocess": None} ev = multiprocessing.Event() self.agent_process = multiprocessing.Process( target=agent.app.run, @@ -139,28 +139,37 @@ def test_root(self): def test_status_write_valid_text(self): """Write a status of 'exception'.""" # First, confirm the status is NOT 'exception'. - _ = self.confirm_status(agent.STATUS_INIT) + _ = self.confirm_status(str(agent.Status.INIT)) form = {"status": "exception"} url_part = "status" _ = self.post_form(url_part, form) - _ = self.confirm_status("exception") + _ = self.confirm_status(str(agent.Status.EXCEPTION)) + + def test_status_write_valid_number(self): + """Write a status of '5'.""" + # First, confirm the status is NOT 'exception'. + _ = self.confirm_status(str(agent.Status.INIT)) + form = {"status": 5} + url_part = "status" + _ = self.post_form(url_part, form) + _ = self.confirm_status(str(agent.Status.EXCEPTION)) def test_status_write_invalid(self): """Fail to provide a valid status.""" form = {"description": "Test Status"} js = self.post_form("status", form, 400) - assert js["message"] == "No status has been provided" + assert js["message"] == "No valid status has been provided" form = {"status": "unexpected value"} - js = self.post_form("status", form, 200) - assert js["message"] == "Analysis status updated" - _ = self.confirm_status("unexpected value") + js = self.post_form("status", form, 400) + assert js["message"] == "No valid status has been provided" + _ = self.confirm_status(str(agent.Status.INIT)) # Write an unexpected random number. form = {"status": random.randint(50, 99)} - js = self.post_form("status", form, 200) - assert js["message"] == "Analysis status updated" - _ = self.confirm_status(str(form["status"])) + js = self.post_form("status", form, 400) + assert js["message"] == "No valid status has been provided" + _ = self.confirm_status(str(agent.Status.INIT)) def test_logs(self): """Test that the agent responds to a request for the logs.""" @@ -447,7 +456,7 @@ def test_execute_py_error_nonexistent_file(self): js = self.post_form("execpy", form, expected_status=200) assert js["message"] == "Successfully executed command" assert "stderr" in js and "No such file or directory" in js["stderr"] - _ = self.confirm_status(agent.STATUS_RUNNING) + _ = self.confirm_status(str(agent.Status.RUNNING)) def test_execute_py_error_non_zero_exit_code(self): """Ensure we get a 400 back when there's a non-zero exit code.""" @@ -463,7 +472,7 @@ def test_execute_py_error_non_zero_exit_code(self): js = self.post_form("execpy", form, expected_status=200) assert js["message"] == "Successfully executed command" assert "hello world" in js["stdout"] - _ = self.confirm_status(agent.STATUS_RUNNING) + _ = self.confirm_status(str(agent.Status.RUNNING)) def test_pinning(self): r = requests.get(f"{BASE_URL}/pinning")