#! /usr/bin/python
### -*-python-*-
###
### Encrypted email address handling
###
### (c) 2006 Mark Wooding
###

###----- Licensing notice ---------------------------------------------------
###
### This program is free software; you can redistribute it and/or modify
### it under the terms of the GNU General Public License as published by
### the Free Software Foundation; either version 2 of the License, or
### (at your option) any later version.
### 
### This program is distributed in the hope that it will be useful,
### but WITHOUT ANY WARRANTY; without even the implied warranty of
### MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
### GNU General Public License for more details.
### 
### You should have received a copy of the GNU General Public License
### along with this program; if not, write to the Free Software Foundation,
### Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

###----- External dependencies ----------------------------------------------

import catacomb as C
import mLib as M
from pysqlite2 import dbapi2 as sqlite
from UserDict import DictMixin
from getopt import getopt, GetoptError
from getdate import getdate 
from sys import stdin, stdout, stderr, exit, argv, exc_info
from email import Parser as EP
import os as OS
import time as T
import sre as RX
import traceback as TB

###----- Database messing ---------------------------------------------------

class AttrDB (object):
  def __init__(me, dbfile):
    me.db = sqlite.connect(dbfile)
  def setup(me):
    cur = me.db.cursor()
    cur.execute('''CREATE TABLE attr
                           (id INTEGER PRIMARY KEY,
                            key VARCHAR(64) NOT NULL,
                            value VARCHAR(256) NOT NULL)''')
    cur.execute('''CREATE TABLE attrset
                           (id INTEGER NOT NULL,
                           attr INTEGER NOT NULL)''')
    cur.execute('''CREATE TABLE uniq
                           (id INTEGER PRIMARY KEY AUTOINCREMENT,
                            dummy INTEGER NOT NULL)''')
    cur.execute('CREATE UNIQUE INDEX attr_bykv ON attr (key, value)')
    cur.execute('CREATE INDEX attrset_byid ON attrset (id)')
    cur.execute('CREATE INDEX attrset_byattr ON attrset (attr)')
    cur.execute('CREATE UNIQUE INDEX attrset_all ON attrset (id, attr)')
  def uniqueid(me):
    cur = me.db.cursor()
    cur.execute('INSERT INTO uniq (dummy) VALUES (0)')
    cur.execute('SELECT MAX(id) FROM uniq')
    id = cur.fetchone()[0]
    cur.execute('DELETE FROM uniq')
    me.commit()
    return id
  def select(me, expr, args = [], cur = None):
    if cur is None: cur = me.db.cursor()
    cur.execute(expr, args)
    while True:
      r = cur.fetchone()
      if r is None: break
      yield r
  def cleanup(me):
    cur = me.db.cursor()
    cur.execute('''DELETE FROM attr WHERE id IN
                           (SELECT attr.id
                            FROM attr LEFT JOIN attrset
                            ON attr.id = attrset.attr
                            WHERE attrset.id ISNULL)''')
  def check(me, cleanp = False):
    toclean = {}
    cur = me.db.cursor()
    for set, attr in me.select('''SELECT attrset.id, attrset.attr
                                          FROM attrset LEFT JOIN attr
                                          ON attrset.attr = attr.id
                                          WHERE attr.id ISNULL''',
                               [], cur):
      print "attrset %d missing attr %d" % (set, attr)
      toclean[set] = True
    if cleanp:
      for set in toclean:
        cur.execute('DELETE FROM attrset WHERE id = ?', [set])
      me.cleanup()
  def commit(me):
    me.db.commit()

class AttrSet (object):
  def __init__(me, db, id = None):
    if id is None: id = db.uniqueid()
    me.id = id
    me.db = db
  def insert(me, key, value):
    cur = me.db.db.cursor()
    try:
      cur.execute('INSERT INTO attr (key, value) VALUES (?, ?)',
                  [key, value])
    except sqlite.OperationalError:
      pass
    cur.execute('SELECT id FROM attr WHERE key = ? AND value = ?',
                [key, value])
    r = cur.fetchone()
    attr = r[0]
    try:
      cur.execute('INSERT INTO attrset VALUES (?, ?)',
                  [me.id, attr])
    except sqlite.OperationalError:
      pass
  def fetch(me):
    for r in me.db.select('''SELECT attr.key, attr.value
                               FROM attr, attrset ON attr.id = attrset.attr
                               WHERE attrset.id = ?''',
                          [me.id]):
      yield r
  def delete(me):
    cur = me.db.db.cursor()
    cur.execute('DELETE FROM attrset WHERE id = ?', [me.id])
    me.db.cleanup()

class AttrMap (AttrSet, DictMixin):
  def __getitem__(me, key):
    it = None
    for v, in me.db.select('''SELECT attr.value
                                FROM attr, attrset ON attr.id = attrset.attr
                                WHERE attrset.id = ? AND attr.key = ?''',
                           [me.id, key]):
      if it is None:
        it = v
      else:
        raise ValueError, 'multiple values for key %s' % key
    if it is None:
      raise KeyError, key
    return it
  def __delitem__(me, key):
    cur = me.db.db.cursor()
    cur.execute('''DELETE FROM attrset
                           WHERE id = ? AND
                                 attr in
                                   (SELECT id FROM attr WHERE key = ?)''',
                [me.id, key])
    me.db.cleanup()
  def __setitem__(me, key, value):
    me.__delitem__(key)
    me.insert(key, value)
  def __iter__(me):
    set = {}
    for k, v in me.fetch():
      if k in set:
        continue
      set[k] = True
      yield k
  def keys(me):
    return [k for k in me]

class AttrMultiMap (AttrMap):
  def __getitem__(me, key):
    them = []
    for v, in me.db.select('''SELECT attr.value
                                FROM attr, attrset ON attr.id = attrset.attr
                                WHERE attrset.id = ? AND attr.key = ?''',
                           [me.id, key]):
      them.append(v)
    if not them:
      raise KeyError, key
    return them
  def __setitem__(me, key, values):
    me.__delitem__(key)
    for it in values:
      me.insert(key, it)

###----- Miscellaneous utilities --------------------------------------------

def time_format(t = None):
  if t is None:
    t = T.time()
  tm = T.gmtime(t)
  return T.strftime('%Y-%m-%d %H:%M:%S', tm)

def any(pred, list):
  for i in list:
    if pred(i): return True
  return False
def every(pred, list):
  for i in list:
    if not pred(i): return False
  return True

prog = RX.sub(r'^.*[/\\]', '', argv[0])
def moan(msg):
  print >>stderr, '%s: %s' % (prog, msg)
def die(msg):
  moan(msg)
  exit(111)

###----- My actual database -------------------------------------------------

class CMDB (AttrDB):
  def setup(me):
    AttrDB.setup(me)
    cur = me.db.cursor()
    cur.execute('''CREATE TABLE expiry
                           (attrset INTEGER PRIMARY KEY,
                            time CHAR(20) NOT NULL)''')
    cur.execute('CREATE INDEX expiry_bytime ON expiry (time)')
  def cleanup(me):
    cur = me.db.cursor()
    now = time_format()
    cur.execute('''DELETE FROM attrset WHERE id IN
                           (SELECT attrset FROM expiry WHERE time < ?)''',
                [now])
    cur.execute('DELETE FROM expiry WHERE time < ?', [now])
    cur.execute('''DELETE FROM expiry WHERE attrset IN
                           (SELECT attrset
                            FROM expiry LEFT JOIN attrset
                            ON expiry.attrset = attrset.id
                            WHERE attrset.id ISNULL)''')
    AttrDB.cleanup(me)
  def expiry(me, id):
    for t, in me.select('SELECT time FROM expiry WHERE attrset = ?', [id]):
      return t
    return None
  def expiredp(me, id):
    t = me.expiry(id)
    if t is not None and t < time_format():
      return True
    else:
      return False
  def setexpire(me, id, when):
    if when != C.KEXP_FOREVER:
      cur = me.db.cursor()
      cur.execute('INSERT INTO expiry VALUES (?, ?)',
                  [id, time_format(when)])

###----- Crypto messing about -----------------------------------------------

## Very vague security arguments...
##
## If the block size n of the PRP is large enough (128 bits) then we encrypt
## id || 0^{n - 64}.  Decryption checks we have the right thing.  The
## security proofs for secrecy and integrity are trivial.
##
## If the block size is small, then we encrypt two blocks:
##   C_0 = E_K(0^{n - 64} || id)
##   C_1 = E_K(C_0)
## The proofs are a little more complicated, but essentially work like this.
## If no 0^{n - 64} || id is ever seen as a C_0 then an adversary can't tell
## the difference between this and a similar construction using independent
## keys.  This other construction must provide secrecy (pushing a
## nonrepeating thing through a PRF) and integrity (PRF on noncolliding
## inputs).  So we win, give or take a birthday term.
class Crypto (object):
  def __init__(me, key):
    me.prp = C.gcprps[key.attr.get('prp', 'blowfish')](key.data.bin)
  def encrypt(me, id):
    blksz = type(me.prp).blksz
    p = C.MP(id).storeb(blksz)
    c = me.prp.encrypt(p)
    if blksz < 16:
      c += me.prp.encrypt(c)
    return c
  def decrypt(me, c):
    bad = False
    blksz = type(me.prp).blksz
    if blksz < 16:
      if len(c) != blksz * 2:
        return None
      c, c1 = c[:blksz], c[blksz:]
      if c1 != me.prp.encrypt(c):
        bad = True
    else:
      if len(c) != blksz:
        return None
    p = me.prp.decrypt(c)
    id = C.MP.loadb(p)
    if id >> 64:
      bad = True
    if bad:
      return None
    return long(id)

###----- Canonification -----------------------------------------------------

rx_prefix = RX.compile(r'''(?x) ^ (
  \[ \S+ \] \s* |
  \S{,4} : \s* |
  \s+
)   
''')
rx_suffix = RX.compile(r'''(?ix) (
  \( \s* was \s* : .* \) \s* |
  \s+
) $''')
rx_punct = RX.compile(r'(?x) [^\w]+ ')

def canon_sender(addr):
  return addr.lower()

def canon_subject(subject):
  subject = subject.lower()
  while True:
    m = rx_prefix.match(subject)
    if not m: break
    subject = subject[m.end():]
  while True:
    m = rx_suffix.search(subject)
    if not m: break
    subject = subject[:m.start()]
  subject = rx_punct.sub('', subject)
  return subject

###----- Checking a message for validity ------------------------------------

class Reject (Exception): pass

class MessageInfo (object):
  __slots__ = '''
    sender msg
  '''.split()

constraints = {}

def check_sender(mi, vv):
  if mi.sender is None:
    raise Reject, 'no sender'
  sender = canon_sender(mi.sender)
  if not any(lambda pat: M.match(pat.lower(), sender), vv):
    raise Reject, 'unmatched sender'
constraints['sender'] = check_sender

def check_subject(mi, vv):
  if mi.msg is None:
    return
  subj = mi.msg['subject']
  if subj is None:
    raise Reject, 'no subject'
  subj = canon_subject(subj)
  if not any(lambda pat: M.match(pat.lower(), subj), vv):
    raise Reject, 'unmatched subject'
constraints['subject'] = check_subject

def check_nothing(me, vv):
  pass
  
def check(db, id, sender = None, msgfile = None):
  mi = MessageInfo()
  a = AttrMultiMap(db, id)
  try:
    addr = a['addr'][0]
  except KeyError:
    raise Reject, 'unknown id'
  if db.expiredp(id):
    raise Reject, 'expired'
  if msgfile is None:
    mi.msg = None
  else:
    try:
      mi.msg = EP.HeaderParser().parse(msgfile)
    except EP.Errors.HeaderParseError:
      raise Reject, 'unparseable header'
  mi.sender = sender
  for k, vv in a.iteritems():
    constraints.get(k, check_nothing)(mi, vv)
  return a['addr'][0]

###----- Commands -----------------------------------------------------------

keyfile = 'db/keyring'
tag = 'cryptomail'
dbfile = 'db/cryptomail.db'
user = None
commands = {}

def timecmp(x, y):
  if x == y:
    return 0
  elif x == C.KEXP_FOREVER or y == C.KEXP_EXPIRE:
    return +1
  elif y == C.KEXP_FOREVER or x == C.KEXP_EXPIRE:
    return +1
  else:
    return cmp(x, y)

def token(c, id):
  return M.base32_encode(c.encrypt(id)).strip('=').lower()

def cmd_generate(argv):
  try:
    opts, argv = getopt(argv, 't:c:f:i:',
                        ['expire=', 'timeout=', 'constraint=',
                         'info=', 'format='])
  except GetoptError:
    return 1
  kr = C.KeyFile(keyfile, C.KOPEN_WRITE)
  k = kr[tag]
  db = CMDB(dbfile)
  map = {}
  expwhen = C.KEXP_FOREVER
  format = '%'
  for o, a in opts:
    if o in ('-t', '--expire', '--timeout'):
      if a == 'forever':
        expwhen = C.KEXP_FOREVER
      else:
        expwhen = getdate(a)
    elif o in ('-c', '--constraint'):
      c, v = a.split('=', 1)
      if c not in constraints:
        die("unknown constraint `%s'", c)
      map.setdefault(c, []).append(v)
    elif o in ('-f', '--format'):
      format = a
    elif o in ('-i', '--info'):
      map['info'] = [a]
    else:
      raise 'Barf!'
  if timecmp(expwhen, k.deltime) > 0:
    k.deltime = expwhen
  if len(argv) != 1:
    return 1
  addr = argv[0]
  a = AttrMultiMap(db)
  a.update(map)
  a['addr'] = [addr]
  if user is not None:
    a['user'] = [user]
  db.setexpire(a.id, expwhen)
  print format.replace('%', token(Crypto(k), a.id))
  db.commit()
  kr.save()
commands['generate'] = \
  (cmd_generate, '[-t TIME] [-c TYPE=VALUE] ADDR', """
Generate a new encrypted email address token forwarding to ADDR.

Subcommand options:
  -t, --timeout=TIME		Address should expire at TIME.
  -c, --constraint=TYPE=VALUE	Apply constraint on the use of the address.
  -f, --format=STRING		Substitute token for `%' in STRING.

Constraint types:
  sender			Envelope sender must match glob pattern.
  subject			Message subject must match glob pattern.""")

def cmd_initdb(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  try:
    OS.unlink(dbfile)
  except OSError:
    pass
  CMDB(dbfile).setup()
commands['initdb'] = \
  (cmd_initdb, '', """
Initialize an attribute database.""")

def getid(local):
  k = C.KeyFile(keyfile, C.KOPEN_READ)[tag]
  id = Crypto(k).decrypt(M.base32_decode(local))
  if id is None:
    raise Reject, 'decrypt failed'
  return id

def cmd_addrcheck(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  local, sender = (lambda addr, sender = None, *hunoz: (addr, sender))(*argv)
  db = CMDB(dbfile)
  try:
    id = getid(local)
    addr = check(db, id, sender)
  except Reject, msg:
    print '-%s' % msg
    return
  print '+%s' % addr
commands['addrcheck'] = \
  (cmd_addrcheck, 'LOCAL [SENDER [IGNORED ...]]', """
Check address token LOCAL, and report `-REASON' for failure or `+ADDR' for
success.""")

def cmd_fwaddr(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  if len(argv) not in (1, 2):
    return 1
  local, sender = (lambda addr, sender = None: (addr, sender))(*argv)
  db = CMDB(dbfile)
  try:
    id = getid(local)
    if id is None:
      raise Reject, 'decrypt failed'
    addr = check(db, id, sender, stdin)
  except Reject, msg:
    print >>stderr, '%s rejected message: %s' % (prog, msg)
    exit(100)
  stdin.seek(0)
  print addr
commands['fwaddr'] = \
  (cmd_fwaddr, 'LOCAL [SENDER]', """
Check address token LOCAL.  On failure, report reason to stderr and exit
111.  On success, write forwarding address to stdout and exit 0.  Expects
the message on standard input, as a seekable file.""")

ignore = {'user': 1, 'addr': 1}
def show(db, a):
  keys = a.keys()
  keys.sort()
  for k in keys:
    if k in ignore:
      continue
    for v in a[k]:
      print '\t%s: %s' % (k, v)
  expwhen = db.expiry(a.id)
  if expwhen:
    print '\texpires: %s' % expwhen
  else:
    print '\tno-expiry'

def cmd_info(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  if len(argv) != 1:
    return 1
  local = argv[0]
  db = CMDB(dbfile)
  try:
    id = getid(local)
    a = AttrMultiMap(db, id)
    if user is not None and user != a.get('user', [None])[0]:
      raise Reject, 'not your token'
    if 'addr' not in a:
      die('unknown token (expired?)')
    print 'addr: %s' % a['addr'][0]
    show(db, a)
  except Reject, msg:
    die('invalid token')
commands['info'] = \
  (cmd_info, 'LOCAL', """
Exaimne the address token LOCAL, and print information about it to standard
output.""")

def cmd_revoke(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  if len(argv) != 1:
    return 1
  local = argv[0]
  db = CMDB(dbfile)
  try:
    id = getid(local)
    a = AttrMultiMap(db, id)
    if user is not None and user != a.get('user', [None])[0]:
      raise Reject, 'not your token'
    if 'addr' not in a:
      die('unknown token (expired?)')
    a.clear()
    db.cleanup()
    db.commit()
  except Reject, msg:
    die('invalid token')
commands['revoke'] = \
  (cmd_revoke, 'LOCAL', """
Revoke the token LOCAL.""")

def cmd_list(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  if argv:
    return 1
  c = Crypto(C.KeyFile(keyfile, C.KOPEN_READ)[tag])
  db = CMDB(dbfile)
  if not user:
    gen = db.select('SELECT DISTINCT id FROM attrset')
  else:
    gen = db.select('''SELECT DISTINCT attrset.id
                               FROM attrset, attr ON attrset.attr = attr.id
                               WHERE attr.key = 'user' AND attr.value = ?''',
                    [user])
  for id, in gen:
    a = AttrMultiMap(db, id)    
    print '%s %s%s' % \
          (token(c, id),
           a.get('addr', '<no-address>')[0],
           (not user and ' [%s]' % a.get('user', ['<no-user>'])[0] or ''))
    show(db, a)
commands['list'] = \
  (cmd_list, '', """
List the user's tokens and information about them.""")

def cmd_cleanup(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  db = CMDB(dbfile)
  db.cleanup()
  cur = db.db.cursor()
  cur.execute('VACUUM')
  db.commit()
commands['cleanup'] = \
  (cmd_cleanup, '', """
Cleans up the attribute database, disposing of old records and compatifying
the file.""")

def cmd_help(argv):
  try:
    opts, argv = getopt(argv, '', [])
  except GetoptError:
    return 1
  if len(argv) == 0:
    cmd = None
  elif len(argv) == 1:
    try:
      cmd = argv[0]
      ci = commands[cmd]
    except KeyError:
      die("unknown command `%s'" % cmd)
  else:
    return 1    
  version()
  print
  if cmd:
    print 'Usage: %s [-OPTIONS] %s %s' % (prog, cmd, ci[1])
    print ci[2]
  else:
    usage(stdout)
    print """
Handle encrypted email addresses.

Help options:
  -h, --help			Show this help text.
  -v, --version			Show version number.
  -u, --usage			Show a usage message.

Global options:
  -d, --database=FILE		Use FILE as the attribute database.
  -k, --keyring=KEYRING		Use KEYRING as the keyring.
  -t, --tag=TAG			Use TAG as the key tag.
  -U, --user=USER		Claim to be USER.
"""
    cmds = commands.keys()
    cmds.sort()
    print 'Subcommands:'
    for c in cmds:
      print '  %s %s' % (c, commands[c][1])
commands['help'] = \
  (cmd_help, '[COMMAND]', """
Show help for subcommand COMMAND.
""")

###----- Main program -------------------------------------------------------

def usage(file):
  print >>file, \
    'Usage: %s [-d FILE] [-k KEYRING] [-t TAG] COMMAND [ARGS...]' % prog
def version():
  print '%s version 1.0.0' % prog
def help():
  cmd_help()  

def main():
  global argv, user, keyfile, dbfile, tag
  try:
    opts, argv = getopt(argv[1:],
                        'hvud:k:t:U:',
                        ['help', 'version', 'usage',
                         'database=', 'keyring=', 'tag=', 'user='])
  except GetoptError:
    usage(stderr)
    exit(111)
  for o, a in opts:
    if o in ('-h', '--help'):
      help()
      exit(0)
    elif o in ('-v', '--version'):
      version()
      exit(0)
    elif o in ('-u', '--usage'):
      usage(stdout)
      exit(0)
    elif o in ('-d', '--database'):
      dbfile = a
    elif o in ('-k', '--keyring'):
      keyfile = a
    elif o in ('-t', '--tag'):
      tag = a
    elif o in ('-U', '--user'):
      user = a
    else:
      raise 'Barf!'
  if len(argv) < 1:
    usage(stderr)
    exit(111)

  if argv[0] in commands:
    c = argv[0]
    argv = argv[1:]
  else:
    usage(stderr)
    exit(111)
  cmd = commands[c]
  if cmd[0](argv):
    print >>stderr, 'Usage: %s %s %s' % (prog, c, cmd[1])
    exit(111)

try:
  main()
except SystemExit:
  raise
except:
  ty, exc, tb = exc_info()
  moan('unhandled %s exception' % ty.__name__)
  for file, line, func, text in TB.extract_tb(tb):
    print >>stderr, \
          '  %-35s -- %.38s' % ('%s:%d (%s)' % (file, line, func), text)
  die('%s: %s' % (ty.__name__, exc[0]))