Skip to content

Commit

Permalink
Unit tests for agent
Browse files Browse the repository at this point in the history
- Added file test_agent.py
- Tests can be run in Windows or Linux
- Tests will be run in github actions
- Test most existing functionality of the agent
- In send_file open the file in binary mode (bug fix)
- Updates to the agent to make it testable, including:
  - Pass a multiprocessing event to the run() method when under test,
    so the test knows when the agent process is ready
  - Tweaks to the shutdown method enabling testing
- Let jsonify not crash if values cannot be serialized
- Add a new command-line parameter, -v, useful when testing interactively
  - When -v is given, stdout and stderr will simply go to the console
- Allow the 'date' command to be executed from localhost; for testing

ran ruff.
  • Loading branch information
rkoumis committed Mar 6, 2024
1 parent 1b03234 commit 93102db
Show file tree
Hide file tree
Showing 4 changed files with 549 additions and 21 deletions.
13 changes: 10 additions & 3 deletions .github/workflows/python-package-windows.yml
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,21 @@ jobs:

- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
# Use x86 python because of https://github.com/kevoreilly/CAPEv2/issues/168
with:
python-version: ${{ matrix.python-version }}
cache: 'pip'
architecture: 'x86'

- name: Install pytest
run: pip install pytest
- name: Install dependencies
run: pip install --upgrade pytest requests

- name: Run unit tests
- name: Run analyzer unit tests
run: |
cd analyzer/windows
pytest -v .
- name: Run agent unit tests
run: |
cd agent
pytest -v .
75 changes: 57 additions & 18 deletions agent/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import http.server
import ipaddress
import json
import multiprocessing
import os
import platform
import shutil
Expand Down Expand Up @@ -36,7 +37,7 @@
if sys.maxsize > 2**32 and sys.platform == "win32":
sys.exit("You should install python3 x86! not x64")

AGENT_VERSION = "0.12"
AGENT_VERSION = "0.13"
AGENT_FEATURES = [
"execpy",
"execute",
Expand All @@ -54,10 +55,6 @@
ANALYZER_FOLDER = ""
state = {"status": STATUS_INIT}

# To send output to stdin comment out this 2 lines
sys.stdout = StringIO()
sys.stderr = StringIO()


class MiniHTTPRequestHandler(http.server.SimpleHTTPRequestHandler):
server_version = "CAPE Agent"
Expand Down Expand Up @@ -105,9 +102,19 @@ def __init__(self):
"POST": [],
}

def run(self, host: ipaddress.IPv4Address = "0.0.0.0", port: int = 8000):
def run(
self,
host: ipaddress.IPv4Address = "0.0.0.0",
port: int = 8000,
event: multiprocessing.Event = None,
):
socketserver.TCPServer.allow_reuse_address = True
self.s = socketserver.TCPServer((host, port), self.handler)
self.s.allow_reuse_address = True

# tell anyone waiting that they're good to go
if event:
event.set()

self.s.serve_forever()

def route(self, path: str, methods: Iterable[str] = ["GET"]):
Expand Down Expand Up @@ -142,24 +149,41 @@ def handle(self, obj):
elif isinstance(ret, send_file):
ret.write(obj.wfile)

if hasattr(self, "s") and self.s._BaseServer__shutdown_request:
self.close_connection = True

def shutdown(self):
# BaseServer also features a .shutdown() method, but you can't use
# that from the same thread as that will deadlock the whole thing.
self.s._BaseServer__shutdown_request = True
if hasattr(self, "s"):
self.s._BaseServer__shutdown_request = True
else:
# When running unit tests in Windows, the system would hang here,
# until this `exit(1)` was added.
print(f"{self} has no 's' attribute")
exit(1)


class jsonify:
"""Wrapper that represents Flask.jsonify functionality."""

def __init__(self, **kwargs):
self.status_code = 200
def __init__(self, status_code=200, **kwargs):
self.status_code = status_code
self.values = kwargs

def init(self):
pass

def json(self):
return json.dumps(self.values)
for valkey in self.values:
if isinstance(self.values[valkey], bytes):
self.values[valkey] = self.values[valkey].decode("utf8", "replace")
try:
retdata = json.dumps(self.values)
except Exception as ex:
retdata = json.dumps({"error": f"Error serializing json data: {ex.args[0]}"})

return retdata

def headers(self, obj):
pass
Expand All @@ -183,7 +207,7 @@ def write(self, sock):
if not self.length:
return

with open(self.path, "r") as f:
with open(self.path, "rb") as f:
buf = f.read(1024 * 1024)
while buf:
sock.write(buf)
Expand Down Expand Up @@ -261,7 +285,13 @@ def put_status():

@app.route("/logs")
def get_logs():
return json_success("Agent logs", stdout=sys.stdout.getvalue(), stderr=sys.stderr.getvalue())
if isinstance(sys.stdout, StringIO):
stdoutbuf = sys.stdout.getvalue()
stderrbuf = sys.stderr.getvalue()
else:
stdoutbuf = "verbose mode, stdout not saved"
stderrbuf = "verbose mode, stderr not saved"
return json_success("Agent logs", stdout=stdoutbuf, stderr=stderrbuf)


@app.route("/system")
Expand Down Expand Up @@ -394,14 +424,17 @@ def do_remove():

@app.route("/execute", methods=["POST"])
def do_execute():
hostname = socket.gethostname()
local_ip = socket.gethostbyname(hostname)
local_ip = socket.gethostbyname(socket.gethostname())

if request.client_ip in ("127.0.0.1", local_ip):
return json_error(500, "Not allowed to execute commands")
if "command" not in request.form:
return json_error(400, "No command has been provided")

# only allow date command from localhost. Even this is just to
# let it be tested
allowed_commands = ["date", "cmd /c date /t"]
if request.client_ip in ("127.0.0.1", local_ip) and request.form["command"] not in allowed_commands:
return json_error(500, "Not allowed to execute commands")

# Execute the command asynchronously? As a shell command?
async_exec = "async" in request.form
shell = "shell" in request.form
Expand Down Expand Up @@ -475,9 +508,15 @@ def do_kill():


if __name__ == "__main__":
multiprocessing.set_start_method("spawn")
parser = argparse.ArgumentParser()
parser.add_argument("host", nargs="?", default="0.0.0.0")
parser.add_argument("port", type=int, nargs="?", default=8000)
# ToDo redir to stdout
parser.add_argument("-v", "--verbose", action="store_true")
args = parser.parse_args()

if not args.verbose:
sys.stdout = StringIO()
sys.stderr = StringIO()

app.run(host=args.host, port=args.port)
3 changes: 3 additions & 0 deletions agent/pytest.ini
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
[pytest]
pythonpath = .
asyncio_mode = auto
Loading

0 comments on commit 93102db

Please sign in to comment.