#!/usr/local/bin/python

# $Id$

# Copyright (C) 2013  Dragon Research Labs ("DRL")
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND DRL DISCLAIMS ALL WARRANTIES WITH
# REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
# AND FITNESS.  IN NO EVENT SHALL DRL 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.

# Preliminary script to work out what's involved in building an
# SQLite3 database of RP objects.  We haven't bothered with this until
# now in rcynic, because we mostly just walk the filesystem tree, but
# LTA and some of the ideas Tim is playing with require a lot of
# lookups based on things that are not the URIs we use as filenames,
# so some kind of indexing may become necessary.  Given the complexity
# of building any kind of real index over RFC 3779 resources,
# otherwise fine lightweight tools like the Python shelve library
# probably won't cut it here, and I don't want to add a dependency on
# MySQL on the RP side (yet?), so let's see what we can do with SQLite3.

import os
import sys
import yaml
import base64
import sqlite3
import weakref
import rpki.POW
import rpki.x509
import rpki.resource_set

sqlite3.register_adapter(rpki.POW.IPAddress,
                         lambda x: buffer("_" + x.toBytes()))

sqlite3.register_converter("RangeVal",
                           lambda s: long(s) if s.isdigit() else rpki.POW.IPAddress.fromBytes(s[1:]))


def main():
  rpdb = RPDB()
  rpdb.load()
  test(rpdb)
  rpdb.close()

def test(rpdb):
  fn2s = [None] + rpdb.fn2map.keys()

  print
  print "Testing YAML parsing"
  parse_yaml(rpdb)

  print
  print "Looking for certificates without AKI"
  for r in rpdb.find_by_aki(None, "cer"):
    print r, r.uris
  print
  print "Testing range functions"
  for fn2 in fn2s:
    if fn2 is not None:
      print
      print "Restricting search to type", fn2
    print
    print "Looking for range that should include adrilankha and psg again"
    for r in rpdb.find_by_range("147.28.0.19", "147.28.0.62", fn2):
      print r, r.uris
    print
    print "Looking for range that should include adrilankha"
    for r in rpdb.find_by_range("147.28.0.19", "147.28.0.19", fn2):
      print r, r.uris
    print
    print "Looking for range that should include ASN 3130"
    for r in rpdb.find_by_range(3130, 3130, fn2):
      print r, r.uris
  print
  print "Moving on to resource sets"
  for fn2 in fn2s:
    if fn2 is not None:
      print
      print "Restricting search to type", fn2
    for expr in ("147.28.0.19-147.28.0.62",
                 "3130",
                 "2001:418:1::19/128",
                 "147.28.0.19-147.28.0.62,198.180.150.50/32",
                 "3130,147.28.0.19-147.28.0.62,198.180.150.50/32",
                 "2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128",
                 "147.28.0.19-147.28.0.62,2001:418:1::19/128,2001:418:1::62/128,198.180.150.50/32,2001:418:8006::50/128"):
      print
      print "Trying", expr
      for r in rpdb.find_by_resource_bag(rpki.resource_set.resource_bag.from_str(expr), fn2):
        print r, r.uris


def parse_xki(s):
  """
  Parse text form of an SKI or AKI.  We accept two encodings:
  colon-delimited hexadecimal, and URL-safe Base64.  The former is
  what OpenSSL prints in its text representation of SKI and AKI
  extensions; the latter is the g(SKI) value that some RPKI CA engines
  (including rpkid) use when constructing filenames.

  In either case, we check that the decoded result contains the right
  number of octets to be a SHA-1 hash.
  """

  if ":" in s:
    b = "".join(chr(int(c, 16)) for c in s.split(":"))
  else:
    b = base64.urlsafe_b64decode(s + ("=" * (4 - len(s) % 4)))  
  if len(b) != 20:
    raise RuntimeError("Bad length for SHA1 xKI value: %r" % s)
  return b


def parse_yaml(rpdb, fn = "rp-sqlite.yaml"):
  yy = yaml.safe_load(open(fn, "r"))
  for y in yy:

    ski = None
    uri = None
    obj = set()

    print

    if "ski" in y:
      ski = parse_xki(y["ski"])
      obj.update(rpdb.find_by_ski(ski))
    if "uri" in y:
      uri = y["uri"]
      obj.update(rpdb.find_by_uri(uri))
    if len(obj) == 1:
      obj = obj.pop()
    else:
      raise RuntimeError("Constraint entry must name a unique object using SKI, URI, or both (%r, %r, %r)" % (
        ski, uri, obj))

    print "URI:", uri
    print "SKI:", " ".join("%02X" % ord(c) for c in ski), "(" + y["ski"] + ")"

    new_resources = old_resources = obj.get_3779resources()

    if "set" in y:
      new_resources = rpki.resource_set.resource_bag.from_str(y["set"])

    if "add" in y:
      new_resources = new_resources | rpki.resource_set.resource_bag.from_str(y["add"]) 

    if "sub" in y:
      new_resources = new_resources - rpki.resource_set.resource_bag.from_str(y["sub"])

    if new_resources == old_resources:
      print "No resource change, skipping"
      continue

    print "Old:", old_resources
    print "New:", new_resources
    print "Add:", new_resources - old_resources
    print "Sub:", old_resources - new_resources

  # See draft-ietf-sidr-ltamgmt-08.txt for real processing details, but overview:
  #
  # - Process constraints file as above to determine list of target
  #   certificates (2.1).  May need to add more fields to YAML hash
  #   for things like CP, CRLDP, etc, although I'm not entirely sure
  #   yet which of those it really makes sense to tweak via
  #   constraints.
  #
  # - Use resources from selected target certificates to determine
  #   which additional certificates we need to reissue to remove those
  #   resources (2.2, "perforation").  In theory we already have SQL
  #   that will just locate all of these for us.
  #
  # - Figure out which trust anchors to process (2.3, TA
  #   re-parenting); we can look in SQL for NULL AKI, but that's just
  #   a hint, we either have to verify that rcynic accepted those TAs
  #   or we have to look at the TALs.  Looking at TALs is probably
  #   easier.
  #
  # At some point we probably need to parse the constraints file into
  # Constraints objects or something like that, except that we may
  # really need something more general that will accomodate
  # perforation and TA reparenting as well.  Figure out and refactor
  # as we go along, most likely.


class RPDB(object):
  """
  Relying party database.

  For now just wire in the database name and rcynic root, fix this
  later if overall approach seems usable.  Might even end up just
  being an in-memory SQL database, who knows?
  """

  fn2map = dict(cer = rpki.x509.X509,
                crl = rpki.x509.CRL,
                mft = rpki.x509.SignedManifest,
                roa = rpki.x509.ROA,
                gbr = rpki.x509.Ghostbuster)

  def __init__(self, db_name = "rp-sqlite.db", delete_old_db = True):

    if delete_old_db:
      try:
        os.unlink(db_name)
      except:
        pass

    exists = os.path.exists(db_name)

    self.db = sqlite3.connect(db_name, detect_types = sqlite3.PARSE_DECLTYPES)
    self.db.text_factory = str
    self.cur = self.db.cursor()

    self.cache = weakref.WeakValueDictionary()

    if exists:
      return

    self.cur.executescript('''
          PRAGMA foreign_keys = on;

          CREATE TABLE object (
                  id INTEGER PRIMARY KEY NOT NULL,
                  der BLOB NOT NULL,
                  fn2 TEXT NOT NULL,
                  ski BLOB,
                  aki BLOB,
                  inherits BOOLEAN NOT NULL,
                  UNIQUE (der));

          CREATE TABLE uri (
                  id INTEGER NOT NULL,
                  uri TEXT NOT NULL,
                  UNIQUE (uri),
                  FOREIGN KEY (id) REFERENCES object(id)
                          ON DELETE CASCADE
                          ON UPDATE CASCADE);

          CREATE INDEX uri_index ON uri(id);

          CREATE TABLE range (
                  id INTEGER NOT NULL,
                  min RangeVal NOT NULL,
                  max RangeVal NOT NULL,
                  UNIQUE (id, min, max),
                  FOREIGN KEY (id) REFERENCES object(id)
                          ON DELETE CASCADE
                          ON UPDATE CASCADE);

          CREATE INDEX range_index ON range(min, max);
          ''')

  def load(self,
           rcynic_root = os.path.expanduser("~/rpki/subvert-rpki.hactrn.net/trunk/"
                                            "rcynic/rcynic-data/unauthenticated"),
           spinner = 100):

    nobj = 0

    for root, dirs, files in os.walk(rcynic_root):
      for fn in files:
        fn = os.path.join(root, fn)
        fn2 = os.path.splitext(fn)[1][1:]

        try:
          obj = self.fn2map[fn2](DER_file = fn)
        except:
          continue

        if spinner and nobj % spinner == 0:
          sys.stderr.write("\r%s %d..." % ("|\\-/"[(nobj/spinner) & 3], nobj))

        nobj += 1

        inherits = False

        if fn2 == "crl":
          ski = None
          aki = buffer(obj.get_AKI())
          cer = None
          bag = None

        else:
          if fn2 == "cer":
            cer = obj
          else:
            cer = rpki.x509.X509(POW = obj.get_POW().certs()[0])
          ski = buffer(cer.get_SKI())
          try:
            aki = buffer(cer.get_AKI())
          except:
            aki = None
          bag = cer.get_3779resources()
          inherits = bag.asn.inherit or bag.v4.inherit or bag.v6.inherit

        der = buffer(obj.get_DER())
        uri = "rsync://" + fn[len(rcynic_root) + 1:]

        try:
          self.cur.execute("INSERT INTO object (der, fn2, ski, aki, inherits) VALUES (?, ?, ?, ?, ?)",
                           (der, fn2, ski, aki, inherits))
          rowid = self.cur.lastrowid

        except sqlite3.IntegrityError:
          self.cur.execute("SELECT id FROM object WHERE der = ? AND fn2 = ?", (der, fn2))
          rows = self.cur.fetchall()
          rowid = rows[0][0]
          assert len(rows) == 1

        else:
          if bag is not None:
            for rset in (bag.asn, bag.v4, bag.v6):
              if rset is not None:
                self.cur.executemany("REPLACE INTO range (id, min, max) VALUES (?, ?, ?)",
                                     ((rowid, i.min, i.max) for i in rset))

        self.cur.execute("INSERT INTO uri (id, uri) VALUES (?, ?)",
                         (rowid, uri))

    if spinner:
      sys.stderr.write("\r= %d objects, committing..." % nobj)

    self.db.commit()

    if spinner:
      sys.stderr.write("done.\n")


  def find_by_ski(self, ski, fn2 = None):
    if ski is None:
      return self._find_results(fn2, "SELECT id, fn2, der FROM object WHERE ski IS NULL")
    else:
      return self._find_results(fn2, "SELECT id, fn2, der FROM object WHERE ski = ?", [buffer(ski)])

  def find_by_aki(self, aki, fn2 = None):
    if aki is None:
      return self._find_results(fn2, "SELECT id, fn2, der FROM object WHERE aki IS NULL")
    else:
      return self._find_results(fn2, "SELECT id, fn2, der FROM object WHERE aki = ?", [buffer(aki)])

  def find_by_uri(self, uri):
    return self._find_results(None, "SELECT object.id, fn2, der FROM object, uri WHERE uri.uri = ? AND object.id = uri.id", [uri])


  # It's easiest to understand overlap conditions by understanding
  # non-overlap then inverting and and applying De Morgan's law.  Ranges
  # A and B do not overlap if either A.min > B.max or A.max < B.min;
  # therefore they do overlap if A.min <= B.max and A.max >= B.min.

  def find_by_range(self, range_min, range_max = None, fn2 = None):
    if range_max is None:
      range_max = range_min
    if isinstance(range_min, (str, unicode)):
      range_min = long(range_min) if range_min.isdigit() else rpki.POW.IPAddress(range_min)
    if isinstance(range_max, (str, unicode)):
      range_max = long(range_max) if range_max.isdigit() else rpki.POW.IPAddress(range_max)
    assert isinstance(range_min, (int, long, rpki.POW.IPAddress))
    assert isinstance(range_max, (int, long, rpki.POW.IPAddress))
    return self._find_results(
      fn2,
      """
      SELECT object.id, fn2, der FROM object, range
      WHERE ? <= max AND ? >= min AND object.id = range.id
      """,
      [range_min, range_max])


  def find_by_resource_bag(self, bag, fn2 = None):
    assert bag.asn or bag.v4 or bag.v6
    qset = []
    aset = []
    for rset in (bag.asn, bag.v4, bag.v6):
      if rset:
        for r in rset:
          qset.append("(? <= max AND ? >= min)")
          aset.append(r.min)
          aset.append(r.max)
    return self._find_results(
      fn2,
      """
      SELECT object.id, fn2, der FROM object, range
      WHERE object.id = range.id AND (%s)
      """ % (" OR ".join(qset)),
      aset)


  def _find_results(self, fn2, query, args = None):
    if args is None:
      args = []
    if fn2 is not None:
      assert fn2 in self.fn2map
      query += " AND fn2 = ?"
      args.append(fn2)
    query += " GROUP BY object.id"
    results = []
    self.cur.execute(query, args)
    selections = self.cur.fetchall()
    for rowid, fn2, der in selections:
      if rowid in self.cache:
        obj = self.cache[rowid]
      else:
        obj = self.fn2map[fn2](DER = der)
        self.cur.execute("SELECT uri FROM uri WHERE id = ?", (rowid,))
        obj.uris = [u[0] for u in self.cur.fetchall()]
        obj.uri = obj.uris[0] if len(obj.uris) == 1 else None
        self.cache[rowid] = obj
      results.append(obj)
    return results


  def close(self):
    self.cur.close()
    self.db.close()


if __name__ == "__main__":
  main()