Stricter checks when setting status

- Use an enum for the status. Only accept expected values.
- Check that we can read a file before trying to send it.

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)
 
 

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")