#!/usr/bin/env python

# $Id$
#
# Copyright (C) 2014  Dragon Research Labs ("DRL")
# Portions copyright (C) 2009-2013  Internet Systems Consortium ("ISC")
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notices and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND DRL AND ISC DISCLAIM ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL DRL OR
# ISC BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL
# DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA
# OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER
# TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
# PERFORMANCE OF THIS SOFTWARE.

import os
import glob
import shutil
import sqlite3
import rpki.rtr.client
import rpki.rtr.channels

def find_free_name(fmt, n = 10):
  """
  Handle SKI hash collisions by allowing a small number of distinct
  filenames for each ASN/SKI pair.
  """

  for i in xrange(n):
    fn = fmt % i
    if not os.path.exists(fn):
      return fn
  raise RuntimeError("Couldn't find a free filename for key %s after %d tries" % (fmt, n))

class ClientChannel(rpki.rtr.client.ClientChannel):
  """
  Subclass ClientChannel to extend .end_of_data() method.
  """

  def end_of_data(self, version, serial, nonce, refresh, retry, expire):
    """
    Call base method to do all the normal EndOfData processing, then
    dump out the key database using the symlink-to-directory hack so
    we can rename() the result to perform an atomic installation.
    """

    # Run the base method
    super(ClientChannel, self).end_of_data(version, serial, nonce, refresh, retry, expire)

    # Set up our new output directory
    dn = "%s.%s" % (self.args.bgpsec_key_directory, rpki.rtr.channels.Timestamp.now())
    ln = "%s.%s" % (self.args.bgpsec_key_directory, ".tmp")
    if os.path.exists(ln):
      os.unlink(ln)
    os.makedirs(dn)

    # Write all the keys
    for asn, gski, key in self.sql.execute("SELECT asn, ski, key FROM routerkey"):
      with open(find_free_name("%s/%s.%s.%%d.key" % (dn, asn, gski)), "wb") as f:
        f.write(key.decode("base64"))

    # Install the new directory
    os.symlink(os.path.basename(dn), ln)
    os.rename(ln, self.args.bgpsec_key_directory)

    # Clean up old output directories
    pattern = ".[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9]T[0-9][0-9]:[0-9][0-9]:[0-9][0-9]Z"
    for gn in glob.iglob(self.args.bgpsec_key_directory + pattern):
      if gn != dn:
        shutil.rmtree(gn)

def SymlinkToDirectory(path):
  """
  argparse "type" function to clean up and sanity check --bgpsec-key-directory.
  """

  path = os.path.abspath(path).rstrip("/")
  if os.path.exists(path) and not os.path.islink(path):
    raise ValueError
  return path

# Grab handle on the normal client argparse setup function
client_argparse_setup = rpki.rtr.client.argparse_setup

# Extend argparse setup to add our own (required) parameter
def argparse_setup(subparsers):
  subparser = client_argparse_setup(subparsers)
  subparser.add_argument("--bgpsec-key-directory", required = True, type = SymlinkToDirectory,
                         help = "where to write BGPSEC router keys")
  return subparser

# Splice our extensions into client
rpki.rtr.client.argparse_setup = argparse_setup
rpki.rtr.client.ClientChannelClass = ClientChannel

# And run the program
import rpki.rtr.main
rpki.rtr.main.main()