generate-hosts.py

#!/usr/bin/env python

from optparse import OptionParser
import os
import re
import datetime


def main():
    usage = "usage: %prog [options] host ..."

    parser = OptionParser(usage, description="Generate hosts file")

    parser.add_option('-s', '--silent', dest='silent', action='store_true',
                      help="echo to stdout", default=False)
    parser.add_option('-w', '--write', dest='write', action='store_true',
                      help="write file back")
    parser.add_option('-u', '--update', dest='update', action='store_true',
                      help="update hosts that already exist")
    parser.add_option('-o', '--overwrite', dest='overwrite', action='store_true',
                      help="overwrite currently existing hosts")
    parser.add_option('-f', '--file', dest='files', action='append',
                      help="read hosts from file", metavar='FILE')
    parser.add_option('-d', '--dir', dest='directories', action='append',
                      help="read folders from target", metavar='DIRECTORY')
    parser.add_option('-t', '--target', dest='target', default='/etc/hosts',
                      help="set hosts target file", metavar='FILE')
    parser.add_option('-a', '--address', dest='address', default='127.0.0.1',
                      help="set default IP address", metavar='IP')

    (options, args) = parser.parse_args()

    # Build initial host dictionary from passed arguments
    hosts = dict([(host, options.address) for host in args])

    # Loop through given directories scanning for subfolders
    if options.directories:
        for directory in options.directories:
            try:
                for item in os.listdir(directory.strip()):
                    if os.path.isdir('%s/%s' % (directory.strip(), item)):
                        hosts[item] = options.address
            except OSError as e:
                print "Unable to list directory '%s'" % directory

    # Loop through given files reading in hosts (+ addresses)
    if options.files:
        for filename in options.files:
            try:
                filehandle = open(filename)
                for line in filehandle:
                    if line and not re.match('^(|\s*#.*)$', line):
                        try:
                            (address, host) = re.search('^\s*(\d+\.\d+\.\d+\.\d+|(?:[a-z0-9]{1,4}:|::){1,7}:?[a-z0-9]{1,4}(?:%\S+)?)\s+(.*)$', line.strip()).groups()
                            if not host in hosts:
                                hosts[host] = address
                        except AttributeError:
                            hosts[line.strip()] = options.address
            except IOError as e:
                print "Unable to read file '%s'" % filename

    try:
        # Read target hosts file, and break into chunks
        current = open(options.target).read() + '\n'
        parts = re.search('\A(.*?)(\n##\n# Local Sites.*?[^\n])\n(\n+.*)?\Z', current, re.DOTALL + re.IGNORECASE)

        if parts:
            (top, sites, bottom) = parts.groups()
        else:
            top = current
            sites = ''
            bottom = ''

        if not bottom:
            bottom = ''

        # Throw out hosts that are already defined in target hosts file
        for line in '\n'.join([top, bottom]).splitlines():
            if line and not re.match('^(|\s*#.*)$', line):
                try:
                    (address, host) = re.search('^\s*(\d+\.\d+\.\d+\.\d+|(?:[a-z0-9]{1,4}:|::){1,7}:?[a-z0-9]{1,4}(?:%\S+)?)\s+(.*)$', line.strip()).groups()
                    if host in hosts:
                        del hosts[host]
                except AttributeError:
                    # Line didn't match address/host search, so naively ignore it
                    pass

        print hosts

        # If we're not overwriting, make sure current defined local sites are added back in
        if not options.overwrite:
            for line in sites.splitlines():
                if line and not re.match('^(|\s*#.*)$', line):
                    try:
                        (address, host) = re.search('^\s*(\d+\.\d+\.\d+\.\d+|(?:[a-z0-9]{1,4}:|::){1,7}:?[a-z0-9]{1,4}(?:%\S+)?)\s+(.*)$', line.strip()).groups()
                        print "%s (%s) - %s" % (host, address, options.update)
                        if not options.update or not host in hosts:
                            hosts[host] = address
                    except AttributeError:
                        # Line didn't match address/host search, so naively ignore it
                        pass

        hosts = [(a, h) for (h, a) in hosts.iteritems()]
        hosts.sort()

        output = []

        # Generate hosts file output
        output.append(top)
        output.append("##")
        output.append("# Local Sites")
        output.append("#")
        output.append("# Automatically generated %s." % datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'))
        output.append("# Change these entries at your peril.")
        output.append("##")

        for host in hosts:
            output.append('%s\t%s' % host)

        output.append(bottom)

        output = '\n'.join(output).strip()

        if not options.silent:
            print output

        if options.write:
            open(options.target, 'w').write(output)

    except IOError as e:
        (code, message) = e.args
        print message


if __name__ == '__main__':
    main()