check.py

#!/usr/bin/env python
"""
You must have nmap installed with the ssl-heartbleed script and its
dependencies (tls.lua) in the current directory.

You also need xsltproc for html output.

The first argument should be a file containing a list of ips, one ip per line,
no blank lines.

Outputs

- nmap xml output to output.xml
- xml output with non-vunlerable hosts stripped to output_parsed.xml
- an html report called report.html
"""

from lxml import etree
import subprocess
import sys

ips = []
with open(sys.argv[1]) as f:
    for line in f.readlines():
        ips.append(line.strip())

cmd = [
    '/usr/local/bin/nmap',
    '-oX',
    'output.xml',
    '-p',
    '443',
    '--script',
    'ssl-heartbleed.nse',
]
cmd.extend(ips)
output = subprocess.check_output(cmd)

with open('output.xml') as file:
    xml = file.read()

root = etree.fromstring(xml)

for host in root.findall('host'):
    script = host.find('.//script')
    if script is None:
        root.remove(host)

vuln_addresses = root.findall('.//address')
vuln_ips = []
for address in vuln_addresses:
    vuln_ips.append(address.get('addr'))

with open('vulnerable_ips.txt', 'w') as f:
    f.write("\n".join(vuln_ips))

parsed_xml = etree.tostring(root, xml_declaration=True)

# stupid, but quick. inserting the stylesheet line into the xml cause
# lxml strips it

out = parsed_xml.split("\n")
stylesheet = '<?xml-stylesheet '
stylesheet += 'href="file:///usr/local/bin/../share/nmap/nmap.xsl" '
stylesheet += 'type="text/xsl"?>'
out.insert(1, stylesheet)
with open('output_parsed.xml', 'w') as file:
    file.write("\n".join(out))

cmd2 = ['xsltproc', 'output_parsed.xml', '--output', 'report.html']

output = subprocess.check_output(cmd2)

XML OUTPUT Header

XML PARSED OUTPUT Header