From 63bfdc40abb781bf01725fe1c98354da4a50dfb9 Mon Sep 17 00:00:00 2001
From: Robin Jadoul <robin.jadoul@gmail.com>
Date: Wed, 24 Mar 2021 20:49:48 +0100
Subject: [PATCH] Initial version, based on challenge 2

---
 __init__.py | 114 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 114 insertions(+)
 create mode 100644 __init__.py

diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..eb05e8b
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,114 @@
+from binaryninja.architecture import Architecture
+from binaryninja.types import Type
+from binaryninja.function import (
+        RegisterInfo,
+        InstructionInfo,
+        InstructionTextToken,
+        InstructionTextTokenType,
+        IntrinsicInfo,
+        IntrinsicInput,
+        )
+
+import struct
+
+def unpack4(x):
+    return struct.unpack("<L", x)[0]
+
+def unpack8(x):
+    return struct.unpack("<Q", x)[0]
+
+class Instr:
+    ops = {
+            0: ("add", "rrr"),
+            1: ("read", "r"),
+            2: ("write", "r"),
+            3: ("movi", "ri"),
+            4: ("xor", "rrr"),
+            9: ("exit", "r"),
+            }
+
+    def __init__(self, data):
+        self.opcode = unpack4(data[:4])
+        self.op, self.argtypes = Instr.ops[self.opcode]
+        self.args = []
+        self.size = 0x1c
+        for i, t in enumerate(self.argtypes):
+            s = 4 + 8 * i
+            self.args.append((t, unpack8(data[s:s+8])))
+
+    def get_text(self):
+        ITT = InstructionTextToken
+        ITTType = InstructionTextTokenType
+        tokens = [ITT(ITTType.InstructionToken, self.op), ITT(ITTType.TextToken, " ")]
+        for t, v in self.args:
+            if t == "r":
+                tokens.append(ITT(ITTType.RegisterToken, f"R{v}"))
+            elif t == "i":
+                tokens.append(ITT(ITTType.IntegerToken, str(v), v))
+            else:
+                assert False
+            tokens.append(ITT(ITTType.OperandSeparatorToken, ", "))
+        return tokens[:-1]
+
+    def get_llil(self, il):
+        args = []
+        for t, v in self.args:
+            if t == "r":
+                args.append(il.reg(8, f"R{v}"))
+            elif t == "i":
+                args.append(il.const(8, v))
+            else:
+                assert False
+        if self.op in ["add", "xor", "movi"]:
+            f = {"add": il.add, "xor": il.xor_expr, "movi": lambda _, a: a}[self.op]
+            v = f(8, *args[1:])
+            return [il.set_reg(8, f"R{self.args[0][1]}", v)]
+        elif self.op == "read":
+            return [il.intrinsic(args, "read", [])]
+        elif self.op == "write":
+            return [il.intrinsic([], "write", args)]
+        elif self.op == "exit":
+            return [il.intrinsic([], "exit", args), il.no_ret()]
+        else:
+            return [il.unimplemented()]
+
+class Poke(Architecture):
+    name = 'poke'
+    instr_alignment = 1
+    max_instr_length = 0x1c
+    address_size = 8
+    default_int_size = 8
+
+    regs = {**{f"R{i}": RegisterInfo(f"R{i}", 8) for i in range(7)},
+                "R7": RegisterInfo("R7", 8),
+                "PC": RegisterInfo("R7", 8),
+                "SP": RegisterInfo("SP", 8),
+            }
+    stack_pointer = "SP"
+
+    intrinsics = {
+                "read": IntrinsicInfo([], [Type.int(1)]),
+                "write": IntrinsicInfo([IntrinsicInput(Type.int(1), "chr")], []),
+                "exit": IntrinsicInfo([IntrinsicInput(Type.int(8), "status")], []),
+                }
+
+    def get_instruction_info(self, data, addr):
+        instr = Instr(data[:self.max_instr_length])
+
+        res = InstructionInfo()
+        res.length = instr.size
+        # No branching yet
+        return res
+
+    def get_instruction_text(self, data, addr):
+        instr = Instr(data[:self.max_instr_length])
+        return instr.get_text(), instr.size
+
+    def get_instruction_low_level_il(self, data, addr, il):
+        instr = Instr(data[:self.max_instr_length])
+        for x in instr.get_llil(il):
+            il.append(x)
+        return instr.size
+
+
+Poke.register()
-- 
2.45.2