GACHAAAAAtkr task solver by gynvael of Dragon Sector

pyvm.py

# GACHAAAAAtkr task solver (potluck ctf, task author: Project Sekai!)
# - by Gynvael Coldwind of Dragon Sector
#
# Note: 99% of this code is python vm reimplementation because I didn't find
# which version of python should I use to run this ;p
# It also uses a timing sidechannel to get the flag.
# Use python 3.12 to run this script!
"""
Dockerfile:

FROM python:3.12.0
WORKDIR /usr/src/app
CMD ["python", "pyvm.py"]

go.sh:

#!/bin/bash
docker build -t my-python-app .
docker run -it --rm --name running-python-app -v "$(pwd)":/usr/src/app my-python-app
"""
import marshal
import types
import opcode
from collections import Counter
import sys
from pprint import pprint

DEBUG = False
DEBUG_VERBOSE = False


_cache_format = { # copied from /Lib/opcode.py
    "LOAD_GLOBAL": {
        "counter": 1,
        "index": 1,
        "module_keys_version": 1,
        "builtin_keys_version": 1,
    },
    "BINARY_OP": {
        "counter": 1,
    },
    "UNPACK_SEQUENCE": {
        "counter": 1,
    },
    "COMPARE_OP": {
        "counter": 1,
    },
    "BINARY_SUBSCR": {
        "counter": 1,
    },
    "FOR_ITER": {
        "counter": 1,
    },
    "LOAD_SUPER_ATTR": {
        "counter": 1,
    },
    "LOAD_ATTR": {
        "counter": 1,
        "version": 2,
        "keys_version": 2,
        "descr": 4,
    },
    "STORE_ATTR": {
        "counter": 1,
        "version": 2,
        "index": 1,
    },
    "CALL": {
        "counter": 1,
        "func_version": 2,
    },
    "STORE_SUBSCR": {
        "counter": 1,
    },
    "SEND": {
        "counter": 1,
    },
}

HANDLERS = {
}



class Null:
  pass
NULL = Null()

class Name:
  def __init__(self, name, value):
    self.n = name
    self.v = value

  def __repr__(self):
    return f"\"{self.n}\" === {self.v}"

  def __str__(self):
    return f"\"{self.n}\" === {self.v}"

def INS_RESUME(ctx, opname, op, arg): pass

def INS_LOAD_NAME(ctx, opname, op, arg):
  n = ctx.c.co_names[arg]
  if DEBUG_VERBOSE:
    print(f"           loading \"{n}\"")
  v = ctx.get_name_value(n)
  ctx.stack.append(Name(n, v))

def INS_LOAD_ATTR(ctx, opname, op, arg):
  lowbit = arg & 1
  namei = arg >> 1
  n = ctx.c.co_names[namei]

  what = ctx.pop()
  if lowbit == 0:
    ctx.push(getattr(what, n))
  else:
    v = getattr(what, n)

    is_bound = getattr(v, "__self__", None) is not None

    if is_bound:
      if DEBUG_VERBOSE: print(f"           bound")
      ctx.push(NULL)
      ctx.push(v)
    else:
      if DEBUG_VERBOSE: print(f"           unbound")
      ctx.push(v)
      ctx.push(what)

def INS_LOAD_CONST(ctx, opname, op, arg):
  ctx.push(ctx.c.co_consts[arg])

def INS_JUMP_BACKWARD(ctx, opname, op, arg):
  ctx.ip = ctx.ip - arg * 2

def INS_JUMP_FORWARD(ctx, opname, op, arg):
  ctx.ip = ctx.ip + arg * 2

def INS_POP_TOP(ctx, opname, op, arg):
  ctx.pop()

def INS_STORE_NAME(ctx, opname, op, arg):
  n = ctx.c.co_names[arg]  # NOTE: I AM NOT IMPLEMENTING GLOBALS HERE YET!
  ctx.set_name_value(n, ctx.pop())

def INS_POP_JUMP_IF_FALSE(ctx, opname, op, arg):
  v = ctx.pop()
  if bool(v) is False:
    ctx.ip += arg * 2
    if DEBUG_VERBOSE: print("           TAKEN")
  else:
    if DEBUG_VERBOSE: print("           not taken")

def INS_POP_JUMP_IF_TRUE(ctx, opname, op, arg):
  v = ctx.pop()
  if bool(v) is True:
    ctx.ip += arg * 2
    if DEBUG_VERBOSE: print("           TAKEN")
  else:
    if DEBUG_VERBOSE: print("           not taken")

def INS_BUILD_TUPLE(ctx, opname, op, arg):
  l = []

  for i in range(arg):
    l.append(ctx.pop())

  # Uff this might be l = l[::-1].
  """
  >>> def a(x,y,z):
  ...   return (x,y,z)
  ...
  >>> dis.dis(a)
    2           0 LOAD_FAST                0 (x)
                2 LOAD_FAST                1 (y)
                4 LOAD_FAST                2 (z)
                6 BUILD_TUPLE              3
                8 RETURN_VALUE
  """
  ctx.push(tuple(l[::-1]))

def INS_BUILD_LIST(ctx, opname, op, arg):
  l = []

  for i in range(arg):
    l.append(ctx.pop())

  # Uff this might be l = l[::-1].
  ctx.push(l[::-1])

def INS_PUSH_NULL(ctx, opname, op, arg):
  ctx.push(NULL)

def INS_CALL(ctx, opname, op, arg):
  global CTX
  CTX = ctx
  args = []
  for _ in range(arg):
    args.append(ctx.pop())
  args = tuple(args[::-1])
  #DEBUG_VERBOSE = True # HAX SHADOWING
  if DEBUG_VERBOSE: print(f"           args: {args}")

  a = ctx.pop()
  b = ctx.pop()

  if b is NULL:
    self = None
    callable = a

    if DEBUG_VERBOSE: print(f"           callable: {callable}")
    if (
      "built-in method read of _io.BufferedReader " not in repr(callable) and
      "built-in method from_bytes of type " not in repr(callable) and
      "built-in method encode of str " not in repr(callable) and
      "built-in method append of bytearray " not in repr(callable) and
      "built-in method pop of bytearray " not in repr(callable) and
      callable not in WHITELIST):
      sys.exit("!!! callable not in WHITELIST!")

    ctx.push(callable(*args))
  else:
    ctx.show_stack()
    self = a
    callable = b
    sys.exit("IMPLEMENT CALLING SOMETHING AS A METHOD")

def INS_CACHE(ctx, opname, op, arg):
  ...  # A bit weird

BINARY_OP = [
    ("NB_ADD", "+"),
    ("NB_AND", "&"),
    ("NB_FLOOR_DIVIDE", "//"),
    ("NB_LSHIFT", "<<"),
    ("NB_MATRIX_MULTIPLY", "@"),
    ("NB_MULTIPLY", "*"),
    ("NB_REMAINDER", "%"),
    ("NB_OR", "|"),
    ("NB_POWER", "**"),
    ("NB_RSHIFT", ">>"),
    ("NB_SUBTRACT", "-"),
    ("NB_TRUE_DIVIDE", "/"),
    ("NB_XOR", "^"),
    ("NB_INPLACE_ADD", "+="),
    ("NB_INPLACE_AND", "&="),
    ("NB_INPLACE_FLOOR_DIVIDE", "//="),
    ("NB_INPLACE_LSHIFT", "<<="),
    ("NB_INPLACE_MATRIX_MULTIPLY", "@="),
    ("NB_INPLACE_MULTIPLY", "*="),
    ("NB_INPLACE_REMAINDER", "%="),
    ("NB_INPLACE_OR", "|="),
    ("NB_INPLACE_POWER", "**="),
    ("NB_INPLACE_RSHIFT", ">>="),
    ("NB_INPLACE_SUBTRACT", "-="),
    ("NB_INPLACE_TRUE_DIVIDE", "/="),
    ("NB_INPLACE_XOR", "^="),
]

"""
This is weird. Why is binary op += when it uses store_fast anyway.
An optimizer hint I guess?
>>> dis.dis(a)
  1           0 RESUME                   0

  2           2 LOAD_FAST_CHECK          0 (g)
              4 LOAD_CONST               1 (1)
              6 BINARY_OP               13 (+=)
             10 STORE_FAST               0 (g)
             12 RETURN_CONST             0 (None)
>>>
"""

def INS_BINARY_OP(ctx, opname, op, arg):
  rhs = ctx.pop()
  lhs_name = ctx.pop(pop_name=True)
  if type(lhs_name) is Name:
    lhs = lhs_name.v
  else:
    lhs = lhs_name


  if DEBUG_VERBOSE: print(f"           {lhs} {BINARY_OP[arg][1]} {rhs}")

  match (arg % 13):
    case 0: ctx.push(lhs + rhs)
    case 1: ctx.push(lhs & rhs)
    case 2: ctx.push(lhs // rhs)
    case 3: ctx.push(lhs << rhs)
    case 4: ctx.push(lhs @ rhs)
    case 5: ctx.push(lhs * rhs)
    case 6: ctx.push(lhs % rhs)
    case 7: ctx.push(lhs | rhs)
    case 8: ctx.push(lhs ** rhs)
    case 9: ctx.push(lhs >> rhs)
    case 10: ctx.push(lhs - rhs)
    case 11: ctx.push(lhs / rhs)
    case 12: ctx.push(lhs ^ rhs)

    # Inplace operations.
    # case 13: lhs += rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 14: lhs &= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 15: lhs //= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 16: lhs <<= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 17: lhs @= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 18: lhs *= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 19: lhs %= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 20: lhs |= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 21: lhs **= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 22: lhs >>= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 23: lhs -= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 24: lhs /= rhs; ctx.set_name_value(lhs_name.v, lhs)
    # case 25: lhs ^= rhs; ctx.set_name_value(lhs_name.v, lhs)
    case _:
      sys.exit("!!! Unknown binary op")

CMP_OP = ('<', '<=', '==', '!=', '>', '>=')

def INS_COMPARE_OP(ctx, opname, op, arg):
  rhs = ctx.pop()
  lhs = ctx.pop()
  cmp_op = arg >> 4  # No, I have no idea what's in the rest of the argument.

  if DEBUG_VERBOSE: print(f"           {lhs} {CMP_OP[cmp_op]} {rhs}")

  match cmp_op:
    case 0: # <
      ctx.push(lhs < rhs)
    case 1: # <=
      ctx.push(lhs <= rhs)
    case 2: # ==
      ctx.push(lhs == rhs)
    case 3: # !=
      ctx.push(lhs != rhs)
    case 4: # >
      ctx.push(lhs > rhs)
    case 5: # >=
      ctx.push(lhs >= rhs)
    case _:
      sys.exit("!!! Unknown compare op")

def INS_BINARY_SLICE(ctx, opname, op, arg):
  end = ctx.pop()
  start = ctx.pop()
  container = ctx.pop()
  ctx.push(container[start:end])

def INS_IMPORT_NAME(ctx, opname, op, arg):
  n = ctx.c.co_names[arg]
  fromlist = ctx.pop()
  level = ctx.pop()

  if n not in {"types"}:
    sys.exit(f"!!! not whitelisted import {n}")

  if DEBUG_VERBOSE: print(f"           importing: {n}, {fromlist}")
  imv = __import__(n, fromlist=fromlist, level=level)
  ctx.push(imv)

def INS_IMPORT_FROM(ctx, opname, op, arg):
  n = ctx.c.co_names[arg]
  m = ctx.stack[-1]

  ctx.push(getattr(m, n))

def INS_EXTENDED_ARG(ctx, opname, op, arg):
  ctx.ext = (ctx.ext << 8) | (arg & 0xff)
  if DEBUG_VERBOSE: print(f"           ext ← 0x{ctx.ext}")

def INS_BINARY_SUBSCR(ctx, opname, op, arg):
  key = ctx.pop()
  container = ctx.pop()
  ctx.push(container[key])

def INS_CONTAINS_OP(ctx, opname, op, arg):
  rhs = ctx.pop()
  lhs = ctx.pop()

  if arg:
    if DEBUG_VERBOSE: print(f"           {lhs} not in {type(rhs)}")
    ctx.push(lhs not in rhs)
  else:
    if DEBUG_VERBOSE: print(f"           {lhs} in {type(rhs)}")
    ctx.push(lhs in rhs)

def INS_STORE_SUBSCR(ctx, opname, op, arg):
  key = ctx.pop()
  container = ctx.pop()
  value = ctx.pop()
  container[key] = value
  if DEBUG_VERBOSE: print(f"           {type(container)}[{key}] ← {value}")

def INS_RETURN_CONST(ctx, opname, op, arg):
  ctx.end = True
  ctx.return_value = ctx.c.co_consts[arg]

def INS_COPY(ctx, opname, op, arg):
  assert arg > 0
  ctx.push(ctx.stack[-arg])

def INS_SWAP(ctx, opname, op, arg):
  i = arg
  ctx.stack[-i], ctx.stack[-1] = ctx.stack[-1], ctx.stack[-i]



# End of instructions.
class Context:
  def __init__(self, c, name, globals):
    self.c = c
    self.name = name
    self.stack = []
    self.ip = 0
    self.oip = None
    self.locals = {}
    self.globals = globals
    self.ext = 0
    self.end = False
    self.return_value = None

  def show_stack(self):
    pprint(self.stack)

  def push(self, v):
    if DEBUG_VERBOSE: print(f"           push {v}")
    self.stack.append(v)

  def pop(self, pop_name=False):
    v = self.stack.pop(-1)
    if DEBUG_VERBOSE:
      print(f"           pop  {v}")
    if pop_name is False and type(v) is Name:
      return v.v
    return v

  def set_name_value(self, name, value):
    #self.locals[name] = value
    self.globals[name] = value
    if DEBUG_VERBOSE: print(f"           {name} ← {value}")


  def get_name_value(self, name):
    if name in self.locals:
      return self.locals[name]

    if name in self.globals:
      return self.globals[name]

    sys.exit(f"Unknown name `{name}`")

def myexec(c, name, globals={}):
  if DEBUG:
    sys.stdout.write("\x1b[1;37m")
    print("=" * 70)
    print(f"→→→ {name}")
    print("=" * 70)
    sys.stdout.write("\x1b[m")



  ctx = Context(c, name, globals)

  d = c._co_code_adaptive  # I *think* this should be adaptive.

  executed = 0
  exec_limit = 0

  while True:
    if (exec_limit != 0 and executed >= exec_limit):
      sys.exit("LIMIT")

    if ctx.end is True:
      if DEBUG:
        print(f"RETURN {ctx.return_value}")
      break


    executed += 1
    global GLOBAL_COUNTER
    GLOBAL_COUNTER += 1

    ctx.oip = ctx.ip
    op = d[ctx.ip]
    arg = d[ctx.ip+1] | (ctx.ext << 8)
    ctx.ip += 2

    opname = opcode.opname[op]
    if opname != "EXTENDED_ARG":
      ctx.ext = 0

    if DEBUG:
      print(f"\x1b[38;2;255;{(ctx.oip*4)&0xff};0m{ctx.name}:{ctx.oip:04}\x1b[m  {opname} {arg}")

    if DEBUG or DEBUG_VERBOSE:
      sys.stdout.write("\x1b[0;32m")

    if opname.startswith("<"):
      sys.exit(f"VERY WRONG OPCODE AT 0x{ctx.oip:x}: {opname}")

    if opname in _cache_format:
      cache_skip = sum(_cache_format[opname].values()) * 2
      if DEBUG_VERBOSE:
        print(f"           <skipping cache {cache_skip} bytes>")
      ctx.ip += cache_skip

    handler_name = f"INS_{opname}"
    handler = HANDLERS.get(handler_name)
    if not handler:
      sys.exit(f"please implement {handler_name}")

    handler(ctx, opname, op, arg)

    if DEBUG or DEBUG_VERBOSE:
      sys.stdout.write("\x1b[m")

  if DEBUG:
    sys.stdout.write("\x1b[1;37m")
    print("-" * 70)
    sys.stdout.write("\x1b[m")
  return ctx


  #pprint(ctx.locals)

def fake_open(name, mode):
  if name == "instructions.bin" and mode == "rb":
    return open("instructions.bin", "rb")

  sys.exit(f"!!! fake_open REJECTED: {name}, {mode}")

def fake_exec(code):
  if CTX.name != "main":
    sys.exit("f!!! REJECTED fake call from not main")

  if "code object main" not in repr(code):
    sys.exit(f"!!! REJECTED not sure what this wanted to execute")

  #CTX.push(True) # lol
  old_ctx = CTX

  #pprint(CTX.locals)
  #sys.exit()
  """
  g = {
    "print": print,
    "exit": sys.exit,
    "bool": bool,
    "IndexError": IndexError,
    "range": range,
    "bytearray": bytearray,
    #"open": fake_open,
    "int": int,
    "list": list,
    #"exec": fake_exec,
    # Giving it access to some locals.
    "D": CTX.locals["D"],
  }
  """
  res_ctx = myexec(code, "innr", CTX.globals)

  return res_ctx.return_value

def fake_exit(*args):
  print(f"!!!fake_exit({args})")
  print(f"GLOBAL_COUNTER: {GLOBAL_COUNTER}")
  sys.exit()

def fake_input():
  #return input()
  return INPUT


def fake_print(*args):
  if "Acc" in args[0]:
    print(*args)
    sys.exit("!!! FLAG")
  #print(*args)
  return

WHITELIST = set([
  bytearray,
  fake_open,
  list,
  types.CodeType,
  fake_exec,
  bool,
  range,
  len,
  fake_input,
  print,
  input,
  fake_print
])



def go(input_text):
  global GLOBAL_COUNTER
  GLOBAL_COUNTER = 0

  global INPUT
  INPUT = input_text

  for k, v in globals().items():
    if k.startswith("INS_"):
      HANDLERS[k] = v

  g = {
    "print": fake_print,
    "exit": sys.exit,
    "bool": bool,
    "IndexError": IndexError,
    "range": range,
    "bytearray": bytearray,
    "open": fake_open,
    "int": int,
    "list": list,
    "exec": fake_exec,
    "len": len,
    "input": fake_input
  }

  with open("mchecker.pyc", "rb") as f:
    d = bytearray(f.read())

  c = marshal.loads(d[16:])
  ctx = myexec(c, "main", g)
  #print("THE END")
  #print(f"{GLOBAL_COUNTER}")
  return GLOBAL_COUNTER

def main():
  # NOTE: you have to run this multiple times to get the flag to "stabilize"
  # It prints position + most likely character based on timing sidechannel.

  for j in range(0,26):

    known = bytearray(b"potluck{.................}")
    res = {}
    for i in range(0x20, 0x7f):
      ch = chr(i)
      known[j] = i
      txt = bytes(known).decode()
      # print(txt, end="  ")
      #print(txt, end="  ")
      cnt = go(txt)
      res[i] = cnt

    mc = Counter(res.values()).most_common()[0][0]

    for k, v in res.items():
      if v == mc:
        continue
      print(j, chr(k))

  """

  for j in range(0, 25):
    known = bytearray(b"..........................")
    res = {}
    for i in range(0x20, 0x7f):
      ch = chr(i)
      known[j] = i
      txt = bytes(known).decode()
      # print(txt, end="  ")
      cnt = go(txt)
      res[i] = cnt

    mc = Counter(res.values()).most_common()[0][0]

    for k, v in res.items():
      if v == mc:
        continue
      print(j, chr(k))
  """





main()
#pprint(opcode.opmap)

# 01234567890123456789012345  ←
# potluck{.................}