bin/expand-zone
author Tero Marttila <terom@paivola.fi>
Thu, 15 Mar 2012 16:08:42 +0200
changeset 1 ea30c9b619b8
child 5 1eb454630f47
permissions -rwxr-xr-x
expand-zone: manage zone serials and expand template vars
#!/usr/bin/env python
# vim: set ft=python :

"""
    Process zonefiles with template expansions.
"""

__version__ = '0.0.1-dev'

import optparse
import codecs
import os.path
from datetime import datetime
import logging

log = logging.getLogger()

# command-line options, global state
options = None

def parse_options (argv) :
    """
        Parse command-line arguments.
    """

    parser = optparse.OptionParser(
            prog        = argv[0],
            usage       = '%prog: [options]',
            version     = __version__,

            # module docstring
            description = __doc__,
    )

    # logging
    general = optparse.OptionGroup(parser, "General Options")

    general.add_option('-q', '--quiet',     dest='loglevel', action='store_const', const=logging.ERROR, help="Less output")
    general.add_option('-v', '--verbose',   dest='loglevel', action='store_const', const=logging.INFO,  help="More output")
    general.add_option('-D', '--debug',     dest='loglevel', action='store_const', const=logging.DEBUG, help="Even more output")

    parser.add_option_group(general)

    parser.add_option('-c', '--input-charset',  metavar='CHARSET',  default='utf-8', 
            help="Encoding used for input files")

    parser.add_option('-o', '--output',         metavar='FILE',     default='-',
            help="Write to output file; default stdout")

    parser.add_option('--output-charset',       metavar='CHARSET',  default='utf-8', 
            help="Encoding used for output files")

    parser.add_option('--expand',               metavar='NAME=VALUE', action='append',
            help="Expand given template variable in zone")

    parser.add_option('--serial',               metavar='FILE',
            help="Read/expand serial from given .serial file")

    parser.add_option('--update-serial',        action='store_true',
            help="Update serial in given .serial file")

    # defaults
    parser.set_defaults(
        loglevel            = logging.WARN,
        expand              = [],
    )
    
    # parse
    options, args = parser.parse_args(argv[1:])

    # configure
    logging.basicConfig(
        format  = '%(processName)s: %(name)s: %(levelname)s %(funcName)s : %(message)s',
        level   = options.loglevel,
    )

    return options, args

def process_file (file, expansions) :
    """
        Process file, expanding lines.
    """

    for line in file :
        line = line.format(**expansions)

        yield line

def write_lines (file, lines, suffix='\n') :
    for line in lines :
        file.write(line + suffix)

def open_file (path, mode, charset) :
    """
        Open unicode-enabled file from path, with - using stdio.
    """

    if path == '-' :
        # use stdin/out based on mode
        stream, func = {
            'r':    (sys.stdin, codecs.getreader),
            'w':    (sys.stdout, codecs.getwriter),
        }[mode[0]]

        # wrap
        return func(charset)(stream)

    else :
        # open
        return codecs.open(path, mode, charset)

def process_serial (path, update=False) :
    """
        Update/process new serial number from given file, based on date.

        Returns the new serial as a string.
    """

    DATE_FMT = '%Y%m%d'
    DATE_LEN = 8

    SERIAL_FMT = "{date:8}{count:02}"
    SERIAL_LEN = 10

    if os.path.exists(path) :
        # read current
        serial = open(path).read().strip()

        assert len(serial) == SERIAL_LEN

        old_serial = int(serial)

        old_date = datetime.strptime(serial[:DATE_LEN], DATE_FMT).date()
        old_count = int(serial[DATE_LEN:])
        
    else :
        log.warn("given .serial does not exist, assuming from today: %s", path)
        old_serial = old_date = old_count = None

    if update :
        # update
        today = datetime.now().date()

        if not old_serial :
            # fresh start
            date = today
            count = 1
            
            log.info("Starting with fresh serial: %s:%s", date, count)
        
        elif old_date < today :
            # update date
            date = today
            count = 1
            
            log.info("Updating to today: %s -> %s", old_date, date)

        elif old_date == today :
            # keep date, update count
            date = old_date
            count = old_count + 1

            if count > 99 :
                raise Exception("Serial update rollover: %s, %s", date, count)

            log.info("Updating today's count: %s, %s", date, count)

        else :
            raise Exception("Invalid serial: %s:%s", old_date, old_count)

    else :
        date = old_date
        count = old_count

    serial = SERIAL_FMT.format(date=date.strftime(DATE_FMT), count=count)

    open(path, 'w').write(serial)

    return serial

def parse_expand (expand) :
    """
        Parse an --expand foo=bar to (key, value)
    """

    key, value = expand.split('=', 1)

    return key, value

def main (argv) :
    global options
    
    options, args = parse_options(argv)

    # expands
    expand = dict(parse_expand(expand) for expand in options.expand)

    # serial?
    if options.serial :
        serial = process_serial(options.serial, update=options.update_serial)

        expand['serial'] = serial

    # input
    if args :
        # open files
        input_files = [open_file(path, 'r', options.input_charset) for path in args]

    else :
        # default to stdout
        input_files = [open_file('-', 'r', options.input_charset)]
   
    # process
    lines = []

    for file in input_files :
        log.info("Reading zone: %s", file)

        lines += list(process_file(file, expand))

    # output
    output = open_file(options.output, 'w', options.output_charset)
    write_lines(output, lines, suffix='')

    return 0

if __name__ == '__main__':
    import sys

    sys.exit(main(sys.argv))