DEFCON CTF 2012 Grab bag 300

grab.py

#!/usr/bin/env python

from __future__ import with_statement
from contextlib import closing
from itertools import imap, ifilter
import socket, re

PW = '5fd78efc6620f6\n'
TARGET = ('140.197.217.85', 10435)
PROMPT = 'Enter ATM PIN:'
ESCAPE_RE = re.compile('\x1b\\[0;[0-9]+;[0-9]+m')
INTERESTING_RE = re.compile(' [0-9] ')
EXAMPLES = 3
DIGITS = 4
INIT_RANGE = range(36)

def main():
    with closing(socket.socket()) as s:
        s.connect(TARGET)
        s.send(PW)
        while True:
            buf = ''
            while PROMPT not in buf:
                r = s.recv(4096)
                buf += r
                if r:
                    print repr(r)
            pin = buffer2pin(buf)
            print 'PIN:', pin
            s.send(pin + '\n')

def buffer2pin(buf):
    buf = ESCAPE_RE.sub('', buf)
    buf = filter(INTERESTING_RE.search, imap(str.strip, buf.split('\n')))
    sets = [set(INIT_RANGE) for _ in xrange(DIGITS)]
    for i in xrange(EXAMPLES):
        base = ''.join(buf[i * 7:i * 7 + 6]).replace(' ', '')
        for n, i in enumerate(ifilter(str.isdigit, buf[i * 7 + 6])):
            sets[n].intersection_update(m.start() for m in re.finditer(i, base))
    quest = ''.join(buf[3 * 7:3 * 7 + 6]).replace(' ', '')
    return ' '.join(quest[digit.pop()] for digit in sets)

if __name__ == '__main__':
    main()