terom@0: #!/usr/bin/env python
terom@0: 
terom@0: """
terom@0:     Process zonefiles.
terom@0: """
terom@0: 
terom@0: __version__ = '0.0.1-dev'
terom@0: 
terom@0: import optparse
terom@0: import codecs
terom@23: from datetime import datetime
terom@0: import logging
terom@0: 
terom@12: log = logging.getLogger('main')
terom@0: 
terom@0: # command-line options, global state
terom@0: options = None
terom@0: 
terom@0: def parse_options (argv) :
terom@0:     """
terom@0:         Parse command-line arguments.
terom@0:     """
terom@0: 
terom@12:     prog = argv[0]
terom@12: 
terom@0:     parser = optparse.OptionParser(
terom@12:             prog        = prog,
terom@0:             usage       = '%prog: [options]',
terom@0:             version     = __version__,
terom@0: 
terom@0:             # module docstring
terom@0:             description = __doc__,
terom@0:     )
terom@0: 
terom@0:     # logging
terom@0:     general = optparse.OptionGroup(parser, "General Options")
terom@0: 
terom@4:     general.add_option('-q', '--quiet',     dest='loglevel', action='store_const', const=logging.ERROR, help="Less output")
terom@0:     general.add_option('-v', '--verbose',   dest='loglevel', action='store_const', const=logging.INFO,  help="More output")
terom@0:     general.add_option('-D', '--debug',     dest='loglevel', action='store_const', const=logging.DEBUG, help="Even more output")
terom@0: 
terom@0:     parser.add_option_group(general)
terom@0: 
terom@12:     # input/output
terom@0:     parser.add_option('-c', '--input-charset',  metavar='CHARSET',  default='utf-8', 
terom@0:             help="Encoding used for input files")
terom@0: 
terom@0:     parser.add_option('-o', '--output',         metavar='FILE',     default='-',
terom@0:             help="Write to output file; default stdout")
terom@0: 
terom@0:     parser.add_option('--output-charset',       metavar='CHARSET',  default='utf-8', 
terom@0:             help="Encoding used for output files")
terom@0: 
terom@23:     # read line mtimes?
terom@23:     parser.add_option('--input-line-date',      action='store_true',
terom@23:             help="Parse timestamp prefix from each input line (e.g. `hg blame | ...`)")
terom@23: 
terom@12:     # check stage
terom@12:     parser.add_option('--check-hosts',          action='store_true',
terom@12:             help="Check that host/IPs are unique. Use --quiet to silence warnings, and test exit status")
terom@12: 
terom@12:     parser.add_option('--check-exempt',         action='append',
terom@12:             help="Allow given names to have multiple records")
terom@12: 
terom@12:     # forward stage
terom@0:     parser.add_option('--forward-zone',         action='store_true', 
terom@0:             help="Generate forward zone")
terom@0: 
terom@0:     parser.add_option('--forward-txt',          action='store_true',
terom@0:             help="Generate TXT records for forward zone")
terom@0: 
terom@0:     parser.add_option('--forward-mx',           metavar='MX',
terom@0:             help="Generate MX records for forward zone")
terom@0: 
terom@12:     # reverse stage
terom@0:     parser.add_option('--reverse-domain',       metavar='DOMAIN',
terom@0:             help="Domain to use for hosts in reverse zone")
terom@0: 
terom@0:     parser.add_option('--reverse-zone',         metavar='NET',
terom@0:             help="Generate forward zone for given subnet (x.z.y)")
terom@0: 
terom@0:     # defaults
terom@0:     parser.set_defaults(
terom@4:         loglevel            = logging.WARN,
terom@12: 
terom@12:         check_exempt        = [],
terom@0:     )
terom@0:     
terom@0:     # parse
terom@0:     options, args = parser.parse_args(argv[1:])
terom@0: 
terom@0:     # configure
terom@0:     logging.basicConfig(
terom@12:         format  = prog + ': %(name)s: %(levelname)s %(funcName)s : %(message)s',
terom@0:         level   = options.loglevel,
terom@0:     )
terom@0: 
terom@0:     return options, args
terom@0: 
terom@23: class ZoneLine (object) :
terom@23:     """
terom@23:         A line in a zonefile.
terom@23:     """
terom@23: 
terom@23:     file = None
terom@23:     lineno = None
terom@23: 
terom@23:     # data
terom@23:     indent = None # was the line indented?
terom@23:     data = None
terom@23:     parts = None # split line fields
terom@23: 
terom@23:     # optional
terom@23:     timestamp = None
terom@23:     comment = None
terom@23: 
terom@23:     PARSE_DATETIME_FORMAT = '%Y-%m-%d'
terom@23: 
terom@23:     @classmethod
terom@23:     def parse (cls, file, lineno, line, line_timestamp_prefix=False) :
terom@23:         """
terom@23:             Parse out given line and build.
terom@23:         """
terom@23: 
terom@23:         log.debug("parse: %s:%d: %s", file, lineno, line)
terom@23: 
terom@23:         ts = None
terom@23: 
terom@23:         if line_timestamp_prefix :
terom@23:             if ': ' not in line :
terom@23:                 raise Exception("Missing timestamp prefix on line: %s:%d: %s" % (file, lineno, line))
terom@23: 
terom@23:             # split prefix
terom@23:             prefix, line = line.split(': ', 1)
terom@23: 
terom@23:             # parse it out
terom@23:             ts = datetime.strptime(prefix, cls.PARSE_DATETIME_FORMAT)
terom@23: 
terom@23:             log.debug("  ts=%r", ts)
terom@23: 
terom@23:         # was line indented?
terom@23:         indent = line.startswith(' ') or line.startswith('\t')
terom@23:         
terom@23:         # strip
terom@23:         line = line.strip()
terom@23:         
terom@23:         log.debug("  indent=%r, line=%r", indent, line)
terom@23: 
terom@23:         # parse comment out?
terom@23:         if ';' in line :
terom@23:             line, comment = line.split(';', 1)
terom@23: 
terom@23:             line = line.strip()
terom@23:             comment = comment.strip()
terom@23: 
terom@23:         else :
terom@23:             line = line.strip()
terom@23:             comment = None
terom@23:         
terom@23:         log.debug("  line=%r, comment=%r", line, comment)
terom@23: 
terom@23:         # parse fields
terom@23:         if '"' in line :
terom@23:             pre, data, post = line.split('"', 2)
terom@23:             parts = pre.split() + [data] + post.split()
terom@23:            
terom@23:         else :
terom@23:             parts = line.split()
terom@23: 
terom@23:         log.debug("  parts=%r", parts)
terom@23: 
terom@23:         # build
terom@23:         return cls(file, lineno, indent, line, parts, timestamp=ts, comment=comment)
terom@23: 
terom@23:     def __init__ (self, file, lineno, indent, data, parts, timestamp=None, comment=None) :
terom@23:         self.file = file
terom@23:         self.lineno = lineno
terom@23: 
terom@23:         self.indent = indent
terom@23:         self.data = data
terom@23:         self.parts = parts
terom@23: 
terom@23:         self.timestamp = timestamp
terom@23:         self.comment = comment
terom@23: 
terom@23:     def __str__ (self) :
terom@23:         return "{file}:{lineno}".format(file=self.file, lineno=self.lineno)
terom@23: 
terom@23: class ZoneRecord (object) :
terom@23:     """
terom@23:         A record from a zonefile.
terom@23:     """
terom@23: 
terom@23:     # the underlying line
terom@23:     line = None
terom@23: 
terom@23:     # record fields
terom@23:     name = None
terom@23:     type = None
terom@23: 
terom@23:     # list of data fields
terom@23:     data = None
terom@23: 
terom@23:     # optional
terom@23:     ttl = None
terom@23:     cls = None
terom@23: 
terom@23:     @classmethod
terom@23:     def parse (cls, line) :
terom@23:         """
terom@23:             Parse from ZoneLine. Returns None if there is no record on the line..
terom@23:         """
terom@23: 
terom@23:         if not line.parts :
terom@23:             # skip
terom@23:             return
terom@23:         
terom@23:         # consume parts
terom@23:         parts = list(line.parts)
terom@23: 
terom@23:         # indented lines don't have name
terom@23:         if line.indent :
terom@23:             name = None
terom@23: 
terom@23:         else :
terom@23:             name = parts.pop(0)
terom@23:         
terom@23:         log.debug("  name=%r", name)
terom@23: 
terom@23:         # parse ttl/cls/type
terom@23:         ttl = _cls = None
terom@23: 
terom@23:         if parts and parts[0][0].isdigit() :
terom@23:             ttl = parts.pop(0)
terom@23: 
terom@23:         if parts and parts[0].upper() in ('IN', 'CH') :
terom@23:             _cls = parts.pop(0)
terom@23: 
terom@23:         # always have type
terom@23:         type = parts.pop(0)
terom@23: 
terom@23:         # remaining parts are data
terom@23:         data = parts
terom@23: 
terom@23:         log.debug("  ttl=%r, cls=%r, type=%r, data=%r", ttl, _cls, type, data)
terom@23: 
terom@23:         return cls(name, type, data,
terom@23:             ttl     = ttl,
terom@23:             cls     = _cls,
terom@23:             line    = line,
terom@23:         )
terom@23: 
terom@23:     def __init__ (self, name, type, data, ttl=None, cls=None, line=None, comment=None) :
terom@23:         self.name = name
terom@23:         self.type = type
terom@23:         self.data = data
terom@23:         
terom@23:         self.ttl = ttl
terom@23:         self.cls = cls
terom@23:         
terom@23:         self.line = line
terom@23: 
terom@23:         # XXX: within line
terom@23:         self._comment = comment
terom@23: 
terom@23:     def build_line (self) :
terom@23:         """
terom@23:             Construct a zonefile-format line..."
terom@23:         """
terom@23: 
terom@23:         # XXX: comment?
terom@23:         if self._comment :
terom@23:             comment = '\t; ' + self._comment
terom@23:         else :
terom@23:             comment = ''
terom@23:             
terom@23:         return u"{name:25} {ttl:4} {cls:2} {type:5} {data}{comment}".format(
terom@23:                 name    = self.name or '',
terom@23:                 ttl     = self.ttl or '',
terom@23:                 cls     = self.cls or '',
terom@23:                 type    = self.type,
terom@23:                 data    = ' '.join(unicode(data) for data in self.data),
terom@23:                 comment = comment,
terom@23:         )
terom@23: 
terom@23:     def __str__ (self) :
terom@23:         return ' '.join((self.name, self.type, ' '.join(self.data)))
terom@23: 
terom@24: class TXTRecord (ZoneRecord) :
terom@24:     """
terom@24:         TXT record.
terom@24:     """
terom@24: 
terom@24:     def __init__ (self, name, text, **opts) :
terom@24:         return super(TXTRecord, self).__init__(name, 'TXT', 
terom@24:             [u'"{0}"'.format(text.replace('"', '\\"'))], 
terom@24:             **opts
terom@24:         )
terom@24: 
terom@23: def parse_record (path, lineno, line, **opts) :
terom@0:     """
terom@0:         Parse (name, ttl, type, data, comment) from bind zonefile.
terom@0: 
terom@0:         Returns None for empty/comment lines.
terom@0:     """
terom@0: 
terom@23:     # line
terom@23:     line = ZoneLine.parse(path, lineno, line, **opts)
terom@23:     record = ZoneRecord.parse(line)
terom@0: 
terom@23:     if record :
terom@23:         return record
terom@4: 
terom@23: def parse_zone_records (file, **opts) :
terom@0:     """
terom@23:         Parse ZoneRecord items from the given zonefile, ignoring non-record lines.
terom@0:     """
terom@0:     
terom@0:     for lineno, line in enumerate(file) :
terom@23:         record = parse_record(file.name, lineno, line, **opts)
terom@0: 
terom@23:         if record :
terom@23:             yield record
terom@0: 
terom@12: def check_zone_hosts (zone, whitelist=None) :
terom@12:     """
terom@12:         Parse host/IP pairs from the zone, and verify that they are unique.
terom@12: 
terom@12:         As an exception, names listed in the given whitelist may have multiple IPs.
terom@12:     """
terom@12: 
terom@12:     by_name = {}
terom@12:     by_ip = {}
terom@12: 
terom@12:     fail = None
terom@12: 
terom@23:     for r in zone :
terom@23:         name = r.name
terom@12: 
terom@12:         # name
terom@12:         if name not in by_name :
terom@23:             by_name[name] = r
terom@12: 
terom@23:         elif r.name in whitelist :
terom@23:             log.debug("Duplicate whitelist entry: %s", r)
terom@12: 
terom@12:         else :
terom@12:             # fail!
terom@23:             log.warn("%s: Duplicate name: %s <-> %s", r.line, r, by_name[name])
terom@12:             fail = True
terom@12: 
terom@12:         # ip
terom@23:         if r.type == 'A' :
terom@23:             ip, = r.data
terom@12: 
terom@12:             if ip not in by_ip :
terom@23:                 by_ip[ip] = r
terom@12: 
terom@12:             else :
terom@12:                 # fail!
terom@23:                 log.warn("%s: Duplicate IP: %s <-> %s", r.line, r, by_ip[ip])
terom@12:                 fail = True
terom@12: 
terom@12:     return fail
terom@12: 
terom@0: def process_zone_forwards (zone, txt=False, mx=False) :
terom@0:     """
terom@0:         Process zone data -> forward zone data.
terom@0:     """
terom@0: 
terom@24:     TIMESTAMP_FORMAT='%Y/%m/%d'
terom@24: 
terom@23:     for r in zone :
terom@23:         yield r
terom@0: 
terom@23:         if r.type == 'A' :
terom@24:             if txt :
terom@24:                 # comment?
terom@24:                 comment = r.line.comment
terom@24: 
terom@24:                 if comment :
terom@24:                     yield TXTRecord(None, comment, ttl=r.ttl)
terom@24: 
terom@24:                 # timestamp?
terom@24:                 timestamp = r.line.timestamp
terom@24: 
terom@24:                 if timestamp :
terom@24:                     yield TXTRecord(None, timestamp.strftime(TIMESTAMP_FORMAT), ttl=r.ttl)
terom@0:             
terom@0:             # XXX: RP, do we need it?
terom@0: 
terom@0:             if mx :
terom@0:                 # XXX: is this a good idea?
terom@23:                 yield ZoneRecord(None, 'MX', [10, mx], ttl=r.ttl)
terom@0: 
terom@0: def reverse_addr (ip) :
terom@0:     """
terom@0:         Return in-addr.arpa reverse for given IPv4 IP.
terom@0:     """
terom@0:     
terom@0:     # parse
terom@0:     octets = tuple(int(part) for part in ip.split('.'))
terom@0: 
terom@0:     for octet in octets :
terom@0:         assert 0 <= octet <= 255
terom@0: 
terom@0:     return '.'.join([str(octet) for octet in reversed(octets)] + ['in-addr', 'arpa'])
terom@0: 
terom@0: def fqdn (*parts) :
terom@0:     return '.'.join(parts) + '.'
terom@0: 
terom@0: 
terom@0: def process_zone_reverse (zone, origin, domain) :
terom@0:     """
terom@0:         Process zone data -> reverse zone data.
terom@0:     """
terom@0: 
terom@23:     for r in zone :
terom@23:         if r.type != 'A' :
terom@0:             continue
terom@0: 
terom@23:         ip, = r.data
terom@0: 
terom@0:         # generate reverse-addr
terom@0:         reverse = reverse_addr(ip)
terom@0: 
terom@0:         # verify
terom@0:         if zone and reverse.endswith(origin) :
terom@0:             reverse = reverse[:-(len(origin) + 1)]
terom@0: 
terom@0:         else :
terom@4:             log.warning("Reverse does not match zone origin, skipping: (%s) -> %s <-> %s", ip, reverse, origin)
terom@0:             continue
terom@0: 
terom@0:         # domain to use
terom@0:         host_domain = domain
terom@23:         host_fqdn = fqdn(r.name, domain)
terom@0: 
terom@23:         yield ZoneRecord(reverse, 'PTR', [host_fqdn])
terom@0: 
terom@23: def write_zone_records (file, zone) :
terom@23:     for r in zone :
terom@23:         file.write(r.build_line() + u'\n')
terom@0: 
terom@0: def open_file (path, mode, charset) :
terom@0:     """
terom@0:         Open unicode-enabled file from path, with - using stdio.
terom@0:     """
terom@0: 
terom@0:     if path == '-' :
terom@0:         # use stdin/out based on mode
terom@0:         stream, func = {
terom@0:             'r':    (sys.stdin, codecs.getreader),
terom@0:             'w':    (sys.stdout, codecs.getwriter),
terom@0:         }[mode[0]]
terom@0: 
terom@0:         # wrap
terom@0:         return func(charset)(stream)
terom@0: 
terom@0:     else :
terom@0:         # open
terom@0:         return codecs.open(path, mode, charset)
terom@0: 
terom@0: def main (argv) :
terom@0:     global options
terom@0:     
terom@0:     options, args = parse_options(argv)
terom@0: 
terom@0:     if args :
terom@0:         # open files
terom@0:         input_files = [open_file(path, 'r', options.input_charset) for path in args]
terom@0: 
terom@0:     else :
terom@0:         # default to stdout
terom@0:         input_files = [open_file('-', 'r', options.input_charset)]
terom@0:    
terom@0:     # process zone data
terom@0:     zone = []
terom@0: 
terom@0:     for file in input_files :
terom@0:         log.info("Reading zone: %s", file)
terom@0: 
terom@23:         zone += list(parse_zone_records(file, 
terom@23:             line_timestamp_prefix   = options.input_line_date,
terom@23:         ))
terom@0: 
terom@12:     # check?
terom@12:     if options.check_hosts :
terom@12:         whitelist = set(options.check_exempt)
terom@12: 
terom@12:         log.debug("checking hosts; whitelist=%r", whitelist)
terom@12: 
terom@12:         if check_zone_hosts(zone, whitelist=whitelist) :
terom@12:             log.warn("Hosts check failed")
terom@12:             return 2
terom@12: 
terom@12:         else :
terom@12:             log.info("Hosts check OK")
terom@12: 
terom@0:     # output file
terom@0:     output = open_file(options.output, 'w', options.output_charset)
terom@0: 
terom@0:     if options.forward_zone :
terom@0:         log.info("Write forward zone: %s", output)
terom@0: 
terom@0:         zone = list(process_zone_forwards(zone, txt=options.forward_txt, mx=options.forward_mx))
terom@0: 
terom@0:     elif options.reverse_zone :
terom@0:         origin = reverse_addr(options.reverse_zone)
terom@0:         domain = options.reverse_domain
terom@0: 
terom@0:         if not domain :
terom@0:             log.error("--reverse-zone requires --reverse-domain")
terom@0:             return 1
terom@0: 
terom@0:         zone = list(process_zone_reverse(zone, origin=origin, domain=domain))
terom@0: 
terom@12:     elif options.check_hosts :
terom@12:         # we only did that, done
terom@12:         return 0
terom@12: 
terom@0:     else :
terom@0:         log.warn("Nothing to do")
terom@12:         return 1
terom@0: 
terom@23:     write_zone_records(output, zone)
terom@0: 
terom@0:     return 0
terom@0: 
terom@0: if __name__ == '__main__':
terom@0:     import sys
terom@0: 
terom@0:     sys.exit(main(sys.argv))