DEFCON CTF [catwestern]

import commands
import socket
import re
import struct
import os

def gen_opcode(s):
    a,b = s.split('=')
    d = {}
    d['rax'] = '\x48\xB8'
    d['rbx'] = '\x48\xBB'
    d['rcx'] = '\x48\xB9'
    d['rdx'] = '\x48\xBA'
    d['rsi'] = '\x48\xBE'
    d['rdi'] = '\x48\xBF'
    d['r8'] = '\x49\xB8'
    d['r9'] = '\x49\xB9'
    d['r10'] = '\x49\xBA'
    d['r11'] = '\x49\xBB'
    d['r12'] = '\x49\xBC'
    d['r13'] = '\x49\xBD'
    d['r14'] = '\x49\xBE'
    d['r15'] = '\x49\xBF'
    return d[a] + struct.pack('<Q', int(b, 16))

def fullhex(s):
    ret = ''
    for c in s:
        b = hex(ord(c))[2:]
        b = b.rjust(2, '0')
        ret += '\\x' + b
    return ret

def conv(r):
    return 'asm("mov %%%%%s,%%0" : "=r"(%s));' % (r, r)

def conv2(s):
    return 'printf("%s=0x%%llx\\n", %s);' % (s,s)

conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
conn.connect(('catwestern_631d7907670909fc4df2defc13f2057c.quals.shallweplayaga.me', 9999))
d = conn.recv(1024)
d = d.split('\n')
d.pop(0)
d = filter(lambda x : x, d)
opcode = ''.join(map(gen_opcode, d))

shellcode = conn.recv(1024)
l = re.search('About to send (\d+) bytes', shellcode, re.DOTALL).group(1)
l = int(l)
shellcode = shellcode[-l:]
shellcode = opcode + shellcode
open('/tmp/shellcode', 'wb').write(shellcode)
shellcode = fullhex(shellcode)

regs = ['rax','rbx','rcx','rdx','rsi','rdi','r8','r9','r10','r11','r12','r13','r14','r15']
asm = map(conv, regs)
asm = '\r\n'.join(asm)

c = map(conv2, regs)
c = '\r\n'.join(c)

s = '''#include <stdio.h>
char shellcode[] = "%s";
int main()
{
long int rax,rbx,rcx,rdx,rsi,rdi,r8,r9,r10,r11,r12,r13,r14,r15;
(*(void (*)()) shellcode)();
%s
%s
}''' % (shellcode, asm, c)

open('/tmp/out.c', 'wb').write(s)
os.system('gcc /tmp/out.c -fno-stack-protector -z execstack -o /tmp/out')
ans =  commands.getoutput('/tmp/out') + '\n'
conn.send(ans)
print conn.recv(1024)
conn.close()