|
| 1 | +import base64 |
| 2 | +import OpenSSL |
| 3 | +import os |
| 4 | +import time |
| 5 | +import fcntl |
| 6 | +import signal |
| 7 | +import tempfile |
| 8 | +import hashlib |
| 9 | +import atexit |
| 10 | +import subprocess |
| 11 | +from datetime import datetime |
| 12 | + |
| 13 | +tmp_path = "/dev/shm/hackergame" |
| 14 | +tmp_flag_path = "/dev/shm" |
| 15 | +conn_interval = int(os.environ["hackergame_conn_interval"]) |
| 16 | +token_timeout = int(os.environ["hackergame_token_timeout"]) |
| 17 | +challenge_timeout = int(os.environ["hackergame_challenge_timeout"]) |
| 18 | +pids_limit = int(os.environ["hackergame_pids_limit"]) |
| 19 | +mem_limit = os.environ["hackergame_mem_limit"] |
| 20 | +flag_path = os.environ["hackergame_flag_path"] |
| 21 | +flag_rule = os.environ["hackergame_flag_rule"] |
| 22 | +challenge_docker_name = os.environ["hackergame_challenge_docker_name"] |
| 23 | +readonly = int(os.environ.get("hackergame_readonly", "1")) |
| 24 | + |
| 25 | +with open("cert.pem") as f: |
| 26 | + cert = OpenSSL.crypto.load_certificate(OpenSSL.crypto.FILETYPE_PEM, f.read()) |
| 27 | + |
| 28 | + |
| 29 | +def validate(token): |
| 30 | + try: |
| 31 | + id, sig = token.split(":", 1) |
| 32 | + sig = base64.b64decode(sig, validate=True) |
| 33 | + OpenSSL.crypto.verify(cert, sig, id.encode(), "sha256") |
| 34 | + return id |
| 35 | + except Exception: |
| 36 | + return None |
| 37 | + |
| 38 | + |
| 39 | +def try_login(id): |
| 40 | + os.makedirs(tmp_path, mode=0o700, exist_ok=True) |
| 41 | + fd = os.open(os.path.join(tmp_path, id), os.O_CREAT | os.O_RDWR) |
| 42 | + fcntl.flock(fd, fcntl.LOCK_EX) |
| 43 | + with os.fdopen(fd, "r+") as f: |
| 44 | + data = f.read() |
| 45 | + now = int(time.time()) |
| 46 | + if data: |
| 47 | + last_login, balance = data.split() |
| 48 | + last_login = int(last_login) |
| 49 | + balance = int(balance) |
| 50 | + last_login_str = ( |
| 51 | + datetime.fromtimestamp(last_login).isoformat().replace("T", " ") |
| 52 | + ) |
| 53 | + balance += now - last_login |
| 54 | + if balance > conn_interval * 3: |
| 55 | + balance = conn_interval * 3 |
| 56 | + else: |
| 57 | + balance = conn_interval * 3 |
| 58 | + if conn_interval > balance: |
| 59 | + print( |
| 60 | + f"Player connection rate limit exceeded, please try again after {conn_interval-balance} seconds. " |
| 61 | + f"连接过于频繁，超出服务器限制，请等待 {conn_interval-balance} 秒后重试。" |
| 62 | + ) |
| 63 | + return False |
| 64 | + balance -= conn_interval |
| 65 | + f.seek(0) |
| 66 | + f.truncate() |
| 67 | + f.write(str(now) + " " + str(balance)) |
| 68 | + return True |
| 69 | + |
| 70 | + |
| 71 | +def check_token(): |
| 72 | + signal.alarm(token_timeout) |
| 73 | + token = input("Please input your token: ").strip() |
| 74 | + id = validate(token) |
| 75 | + if not id: |
| 76 | + print("Invalid token") |
| 77 | + exit(-1) |
| 78 | + if not try_login(id): |
| 79 | + exit(-1) |
| 80 | + signal.alarm(0) |
| 81 | + return token, id |
| 82 | + |
| 83 | + |
| 84 | +def generate_flags(token): |
| 85 | + functions = {} |
| 86 | + for method in "md5", "sha1", "sha256": |
| 87 | + |
| 88 | + def f(s, method=method): |
| 89 | + return getattr(hashlib, method)(s.encode()).hexdigest() |
| 90 | + |
| 91 | + functions[method] = f |
| 92 | + |
| 93 | + if flag_path: |
| 94 | + flag = eval(flag_rule, functions, {"token": token}) |
| 95 | + if isinstance(flag, tuple): |
| 96 | + return dict(zip(flag_path.split(","), flag)) |
| 97 | + else: |
| 98 | + return {flag_path: flag} |
| 99 | + else: |
| 100 | + return {} |
| 101 | + |
| 102 | + |
| 103 | +def generate_flag_files(flags): |
| 104 | + flag_files = {} |
| 105 | + for flag_path, flag in flags.items(): |
| 106 | + with tempfile.NamedTemporaryFile("w", delete=False, dir=tmp_flag_path) as f: |
| 107 | + f.write(flag + "\n") |
| 108 | + fn = f.name |
| 109 | + os.chmod(fn, 0o444) |
| 110 | + flag_files[flag_path] = fn |
| 111 | + return flag_files |
| 112 | + |
| 113 | + |
| 114 | +def cleanup(): |
| 115 | + if child_docker_id: |
| 116 | + subprocess.run( |
| 117 | + f"docker rm -f {child_docker_id}", |
| 118 | + shell=True, |
| 119 | + stdout=subprocess.DEVNULL, |
| 120 | + stderr=subprocess.DEVNULL, |
| 121 | + ) |
| 122 | + for file in flag_files.values(): |
| 123 | + os.unlink(file) |
| 124 | + |
| 125 | + |
| 126 | +def create_docker(flag_files, id): |
| 127 | + cmd = ( |
| 128 | + f"docker create --init --rm -i --network none " |
| 129 | + f"--pids-limit {pids_limit} -m {mem_limit} --memory-swap -1 --cpus 1 " |
| 130 | + f"-e hackergame_token=$hackergame_token " |
| 131 | + ) |
| 132 | + |
| 133 | + if readonly: |
| 134 | + cmd += "--read-only " |
| 135 | + |
| 136 | + if challenge_docker_name.endswith("_challenge"): |
| 137 | + name_prefix = challenge_docker_name[:-10] |
| 138 | + else: |
| 139 | + name_prefix = challenge_docker_name |
| 140 | + |
| 141 | + timestr = datetime.now().strftime("%m%d_%H%M%S_%f")[:-3] |
| 142 | + child_docker_name = f"{name_prefix}_u{id}_{timestr}" |
| 143 | + cmd += f'--name "{child_docker_name}" ' |
| 144 | + |
| 145 | + with open("/proc/self/cgroup") as f: |
| 146 | + for line in f: |
| 147 | + if "/docker/" in line: |
| 148 | + docker_id = line.strip()[-64:] |
| 149 | + break |
| 150 | + prefix = f"/var/lib/docker/containers/{docker_id}/mounts/shm/" |
| 151 | + |
| 152 | + for flag_path, fn in flag_files.items(): |
| 153 | + flag_src_path = prefix + fn.split("/")[-1] |
| 154 | + cmd += f"-v {flag_src_path}:{flag_path}:ro " |
| 155 | + |
| 156 | + cmd += challenge_docker_name |
| 157 | + |
| 158 | + return subprocess.check_output(cmd, shell=True).decode().strip() |
| 159 | + |
| 160 | + |
| 161 | +def run_docker(child_docker_id): |
| 162 | + cmd = f"timeout -s 9 {challenge_timeout} docker start -i {child_docker_id}" |
| 163 | + subprocess.run(cmd, shell=True) |
| 164 | + |
| 165 | + |
| 166 | +if __name__ == "__main__": |
| 167 | + child_docker_id = None |
| 168 | + flag_files = {} |
| 169 | + atexit.register(cleanup) |
| 170 | + |
| 171 | + token, id = check_token() |
| 172 | + os.environ["hackergame_token"] = token |
| 173 | + flags = generate_flags(token) |
| 174 | + flag_files = generate_flag_files(flags) |
| 175 | + child_docker_id = create_docker(flag_files, id) |
| 176 | + run_docker(child_docker_id) |
0 commit comments