Merge pull request #3 from C4T-BuT-S4D/pomo/service-ark

Ark service

Author: pomo-mondreganto

.gitignore

@@ -74,4 +74,4 @@ dmypy.json
 .DS_Store
 .vscode
 .idea
-
+.mise.toml

check.py

@@ -48,6 +48,8 @@
     "sysctls",
     "privileged",
     "security_opt",
+    "ulimits",
+    "command",
 ]
 SERVICE_REQUIRED_OPTIONS = ["pids_limit", "mem_limit", "cpus"]
 SERVICE_ALLOWED_OPTIONS = CONTAINER_ALLOWED_OPTIONS
@@ -213,8 +215,10 @@ def put(self, flag: str, flag_id: str, vuln: int):
         cmd = [str(self._exe_path), "put", HOST, flag_id, flag, str(vuln)]
         out, err = self._run_command(cmd)
 
-        self._fatal(len(out) <= 1024, "returned stdout is longer than 1024 characters")
-        self._fatal(len(err) <= 1024, "returned stderr is longer than 1024 characters")
+        self._fatal(len(out) <= 1024,
+                    "returned stdout is longer than 1024 characters")
+        self._fatal(len(err) <= 1024,
+                    "returned stderr is longer than 1024 characters")
 
         if self._attack_data:
             self._fatal(out, "stdout is empty")
@@ -241,7 +245,8 @@ def run_all(self, step: int):
 
         for vuln in range(1, self._vulns + 1):
             flag = generate_flag(self._name)
-            flag_id = self.put(flag=flag, flag_id=secrets.token_hex(16), vuln=vuln)
+            flag_id = self.put(
+                flag=flag, flag_id=secrets.token_hex(16), vuln=vuln)
             flag_id = flag_id.strip()
             self.get(flag, flag_id, vuln)
 
@@ -330,9 +335,11 @@ def validate_file(self, f: Path):
         path = f.relative_to(BASE_DIR)
 
         if f.name not in ALLOWED_YAML_FILES:
-            self._error(f.suffix != ".yaml", f"file {path} has .yaml extension")
+            self._error(f.suffix != ".yaml",
+                        f"file {path} has .yaml extension")
 
-        self._error(f.name != ".gitkeep", f"{path} found, should be named .keep")
+        self._error(f.name != ".gitkeep",
+                    f"{path} found, should be named .keep")
 
         if f.name == "docker-compose.yml":
             with f.open() as file:
@@ -355,7 +362,8 @@ def validate_file(self, f: Path):
                 try:
                     dc_version = float(dc["version"])
                 except ValueError:
-                    self._error(False, f"version option in {path} is not float")
+                    self._error(
+                        False, f"version option in {path} is not float")
                     return
 
                 self._error(
@@ -426,12 +434,14 @@ def validate_file(self, f: Path):
                     else:
                         context = build["context"]
                         if "dockerfile" in build:
-                            dockerfile = f.parent / context / build["dockerfile"]
+                            dockerfile = f.parent / \
+                                context / build["dockerfile"]
                         else:
                             dockerfile = f.parent / context / "Dockerfile"
 
                     if self._error(
-                        dockerfile.exists(), f"no dockerfile found in {dockerfile}"
+                        dockerfile.exists(
+                        ), f"no dockerfile found in {dockerfile}"
                     ):
                         continue
 
@@ -496,7 +506,8 @@ def validate_file(self, f: Path):
             for p in ALLOWED_CHECKER_PATTERNS:
                 checker_code = checker_code.replace(p, "")
             for p in FORBIDDEN_CHECKER_PATTERNS:
-                self._error(p not in checker_code, f'forbidden pattern "{p}" in {path}')
+                self._error(p not in checker_code,
+                            f'forbidden pattern "{p}" in {path}')
 
     def __str__(self):
         return f"Structure validator for {self._service.name}"

checkers/ark/ark_lib.py

@@ -0,0 +1,127 @@
+import random
+import re
+import string
+from dataclasses import dataclass
+
+from checklib import *  # type: ignore
+from pwn import context, remote
+
+context.timeout = 10
+context.log_level = 'FATAL'
+
+PORT = 13345
+
+
+@dataclass
+class SuggestedUser:
+    username: str
+    file_count: int
+    file_size: int
+
+
+@dataclass
+class UserFile:
+    path: str
+    size: int
+
+
+@dataclass
+class File:
+    path: str
+    size: int
+
+
+class CheckMachine:
+    def __init__(self, checker: BaseChecker):
+        self.c = checker
+        self.port = PORT
+
+    def connect(self) -> remote:
+        r = remote(self.c.host, self.port)
+        r.recvuntil(b'quit')
+        r.recvuntil(b'> ')
+        return r
+
+    def register(self, r: remote, username: str, password: str, status: Status = Status.MUMBLE):
+        r.sendline(f'register {username} {password}'.encode())
+        response = r.recvline().decode()
+        r.recvuntil(b'> ')
+
+        self.c.assert_in('Registration successful', response,
+                         'Registration failed', status=status)
+
+    def login(self, r: remote, username: str, password: str, status: Status = Status.MUMBLE):
+        r.sendline(f'login {username} {password}'.encode())
+        response = r.recvline().decode()
+        r.recvuntil(b'> ')
+
+        self.c.assert_in('Welcome,', response,
+                         'Login failed', status=status)
+
+        return r
+
+    def save_file(self, r: remote, path: str, content: str, status: Status = Status.MUMBLE):
+        r.sendline(f'save {path} {content}'.encode())
+        response = r.recvuntil(b'> ').decode()
+
+        self.c.assert_in('File saved successfully', response,
+                         'Save failed', status=status)
+
+    def cat_file(self, r: remote, path: str, status: Status = Status.MUMBLE) -> str:
+        r.sendline(f'cat {path}'.encode())
+        response = r.recvuntil(b'> ').decode()
+
+        # Extract content between the command and the next prompt
+        lines = response.split('\n')
+        self.c.assert_gt(len(lines), 1, 'Cat failed', status=status)
+        return lines[0]
+
+    def suggest_users(self, r: remote, status: Status = Status.MUMBLE) -> list[SuggestedUser]:
+        r.sendline(b'suggest_users')
+        response = r.recvuntil(b'> ').decode()
+
+        self.c.assert_in('Suggested users:', response,
+                         'Suggest users failed', status=status)
+
+        users = []
+        for line in response.split('\n'):
+            if 'Username:' in line:
+                # Username: kek, Created At: 2024-12-15 13:10:54.753883 UTC, File Count: 1, Total File Size: 4
+                matches = re.match(
+                    r'Username: (\w+), .*, File Count: (\d+), Total File Size: (\d+)', line)
+                if matches:
+                    users.append(SuggestedUser(
+                        matches.group(1), int(matches.group(2)), int(matches.group(3))))
+        return users
+
+    def list_user_files(self, r: remote, username: str, status: Status = Status.MUMBLE) -> list[UserFile]:
+        if username == '':
+            r.sendline(b'list')
+        else:
+            r.sendline(f'list_files {username}'.encode())
+
+        response = r.recvuntil(b'> ').decode()
+
+        files = []
+        for line in response.split('\n'):
+            if 'ID' in line:
+                # ID: 2, Path: /tmp/1, Size: 4 bytes
+                matches = re.match(
+                    r'ID: (\d+), Path: (\S+), Size: (\d+) bytes', line)
+                if matches:
+                    files.append(
+                        UserFile(matches.group(2), int(matches.group(3))))
+        return files
+
+    def copy_file(self, r: remote, src_path: str, dst_path: str, status: Status = Status.MUMBLE):
+        r.sendline(f'copy {src_path} {dst_path}'.encode())
+        response = r.recvuntil(b'> ').decode()
+
+        self.c.assert_in('File copied successfully', response,
+                         'Copy failed', status=status)
+
+    def random_filename(self) -> str:
+        d = random.choice(['tmp', 'files'])
+        l = random.randint(10, 50)
+        a = string.ascii_letters + string.digits + '.,-_=+:|()*[]&^'
+        return f"/{d}/test_{rnd_string(l, a)}.txt"

checkers/ark/checker.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python3
+
+import json
+import random
+import sys
+
+if True:
+    saved_args = sys.argv.copy()
+
+from ark_lib import *
+from checklib import *  # type: ignore
+from pwn import PwnlibException
+
+
+class Checker(BaseChecker):
+    vulns: int = 1
+    timeout: int = 20
+    uses_attack_data: bool = True
+
+    def __init__(self, *args, **kwargs):
+        super(Checker, self).__init__(*args, **kwargs)
+        self.mch = CheckMachine(self)
+
+    def check(self):
+        # Register a user
+        username, password = rnd_username(), rnd_password()
+        with (
+            self.mch.connect() as r1,
+            self.mch.connect() as r2,
+            self.mch.connect() as anon_r,
+        ):
+            def gr():
+                return random.choice([r1, r2])
+
+            def ar():
+                return random.choice([r1, anon_r])
+
+            self.mch.register(gr(), username, password)
+
+            self.mch.login(r1, username, password)
+            self.mch.login(r2, username, password)
+
+            files = [
+                (self.mch.random_filename(), rnd_string(random.randint(10, 50)))
+                for _ in range(random.randint(1, 10))
+            ]
+            for file, content in files:
+                self.mch.save_file(gr(), file, content)
+
+            random.shuffle(files)
+            for file, content in files:
+                got_content = self.mch.cat_file(gr(), file, Status.MUMBLE)
+                self.assert_eq(got_content, content, "File content mismatch")
+
+            # Test list functionality
+            list_output = self.mch.list_user_files(gr(), '')
+            for file, _ in files:
+                self.assert_in(file, [f.path for f in list_output],
+                               "File not found in list")
+
+            # Test suggest_users
+            suggest_output = self.mch.suggest_users(gr())
+            if len(suggest_output) < 30 and not any(u.username == username for u in suggest_output):
+                self.cquit(Status.MUMBLE, "Suggest users failed")
+
+            # Test list_user_files
+            user_files = self.mch.list_user_files(ar(), username)
+            for file, content in files:
+                self.assert_in(file, [f.path for f in user_files],
+                               "File not found in user files")
+                self.assert_in(len(content), [f.size for f in user_files],
+                               "File size mismatch")
+
+            # Register a second user
+            # Test second user can copy first user's files
+            # Test first user can cat their own files from the second user's account
+
+            username2, password2 = rnd_username(), rnd_password()
+            self.mch.register(ar(), username2, password2)
+            self.mch.login(r2, username2, password2)
+
+            # Test list_user_files
+            user_files = self.mch.list_user_files(ar(), username)
+            for file, content in files:
+                self.assert_in(file, [f.path for f in user_files],
+                               "File not found in user files")
+                self.assert_in(len(content), [f.size for f in user_files],
+                               "File size mismatch")
+
+            file_to_copy = random.choice(files)
+            dst_path = self.mch.random_filename()
+            self.mch.copy_file(r2, file_to_copy[0], dst_path)
+
+            # Test second user can list their own files
+            user_files = self.mch.list_user_files(r2, username2)
+            self.assert_in(dst_path, [f.path for f in user_files],
+                           "File not found in user files")
+
+            # Test first user can cat their own files from the second user's quota
+            content = self.mch.cat_file(r1, dst_path, Status.MUMBLE)
+            self.assert_eq(content, file_to_copy[1], "File content mismatch")
+
+            self.cquit(Status.OK)
+
+    def put(self, flag_id: str, flag: str, vuln: str):
+        username1, password1 = rnd_username(), rnd_password()
+        username2, password2 = rnd_username(), rnd_password()
+
+        with (
+            self.mch.connect() as r1,
+            self.mch.connect() as r2,
+        ):
+            self.mch.register(r1, username1, password1)
+            self.mch.register(r2, username2, password2)
+
+            self.mch.login(r1, username1, password1)
+            self.mch.login(r2, username2, password2)
+
+            flag_file1 = self.mch.random_filename()
+            flag_file2 = self.mch.random_filename()
+
+            self.mch.save_file(r1, flag_file1, flag)
+            self.mch.copy_file(r2, flag_file1, flag_file2)
+
+            public = f'u1={username1}:u2={username2}'
+
+            private = json.dumps({
+                'username1': username1,
+                'password1': password1,
+                'username2': username2,
+                'password2': password2,
+                'flag_file1': flag_file1,
+                'flag_file2': flag_file2,
+            })
+
+            self.cquit(Status.OK, public=public, private=private)
+
+    def get(self, flag_id: str, flag: str, vuln: str):
+        data = json.loads(flag_id)
+
+        with (
+            self.mch.connect() as r1,
+            self.mch.connect() as r2,
+            self.mch.connect() as anon_r,
+        ):
+            self.mch.login(r1, data['username1'],
+                           data['password1'], Status.CORRUPT)
+            self.mch.login(r2, data['username2'],
+                           data['password2'], Status.CORRUPT)
+
+            for r in [r1, r2, anon_r]:
+                file_list = self.mch.list_user_files(r, data['username1'])
+                self.assert_in(data['flag_file1'], [f.path for f in file_list],
+                               "Flag file not found in user files", Status.CORRUPT)
+
+                file_list = self.mch.list_user_files(r, data['username2'])
+                self.assert_in(data['flag_file2'], [f.path for f in file_list],
+                               "Flag file not found in user files", Status.CORRUPT)
+
+            content = self.mch.cat_file(r1, data['flag_file1'], Status.CORRUPT)
+            self.assert_eq(content, flag, "Flag mismatch", Status.CORRUPT)
+
+            content = self.mch.cat_file(r1, data['flag_file2'], Status.CORRUPT)
+            self.assert_eq(content, flag, "Flag mismatch", Status.CORRUPT)
+
+        self.cquit(Status.OK)
+
+    def action(self, action, *args, **kwargs):
+        try:
+            super(Checker, self).action(action, *args, **kwargs)
+        except PwnlibException:
+            self.cquit(Status.DOWN, 'Connection error', 'Got pwnlib exception')
+
+
+if __name__ == '__main__':
+    c = Checker(saved_args[2])
+
+    try:
+        c.action(saved_args[1], *saved_args[3:])
+    except c.get_check_finished_exception():
+        cquit(Status(c.status), c.public, c.private)