#!/usr/bin/env python
"""
Polls a mailbox for new registrations and forwards them using fp-reg-decrypt.
"""

import calendar
import datetime
import email
import email.utils
import getopt
import imaplib
import math
import os
import re
import socket
import ssl
import stat
import sys
import tempfile
import time

from flashproxy import fac
from flashproxy import keys
from flashproxy import proc
from flashproxy.util import parse_addr_spec

from hashlib import sha1
from M2Crypto import SSL

# TODO(infinity0): we only support gmail so this is OK for now. in the future,
# could maybe do an MX lookup and try to guess the imap server from that.
DEFAULT_IMAP_HOST = "imap.gmail.com"
DEFAULT_IMAP_PORT = 993
DEFAULT_LOG_FILENAME = "fp-registrar-email.log"

POLL_INTERVAL = 60
# Ignore message older than this many seconds old, or newer than this many
# seconds in the future.
REGISTRATION_AGE_LIMIT = 30 * 60

FACILITATOR_ADDR = ("127.0.0.1", 9002)

LOG_DATE_FORMAT = "%Y-%m-%d %H:%M:%S"

class options(object):
    password_filename = None
    log_filename = DEFAULT_LOG_FILENAME
    log_file = sys.stdout
    daemonize = True
    pid_filename = None
    privdrop_username = None
    safe_logging = True
    imaplib_debug = False
    use_certificate_pin = True

# Like socket.create_connection in that it tries resolving different address
# families, but doesn't connect the socket.
def create_socket(address, timeout = None, source_address = None):
    host, port = address
    addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM)
    if not addrs:
        raise socket.error("getaddrinfo returns an empty list")
    err = None
    for addr in addrs:
        try:
            s = socket.socket(addr[0], addr[1], addr[2])
            if timeout is not None and type(timeout) == float:
                s.settimeout(timeout)
            if source_address is not None:
                s.bind(source_address)
            return s
        except Exception, e:
            err = e
    raise err

class IMAP4_SSL_REQUIRED(imaplib.IMAP4_SSL):
    """A subclass of of IMAP4_SSL that uses ssl_version=ssl.PROTOCOL_TLSv1 and
    cert_reqs=ssl.CERT_REQUIRED."""
    def open(self, host = "", port = imaplib.IMAP4_SSL_PORT):
        ctx = SSL.Context("tlsv1")
        ctx.set_verify(SSL.verify_peer, 3)
        ret = ctx.load_verify_locations(self.certfile)
        assert ret == 1

        self.host = host
        self.port = port
        self.sock = create_socket((self.host, self.port))

        self.sslobj = SSL.Connection(ctx, self.sock)
        self.sslobj.connect((self.host, self.port))
        self.file = self.sslobj.makefile('rb')

def usage(f = sys.stdout):
    print >> f, """\
Usage: %(progname)s --pass=PASSFILE
Facilitator-side helper for the flashproxy-reg-email rendezvous. Polls
an IMAP server for email messages with client registrations, deletes
them, and forwards the registrations to the facilitator.

  -d, --debug               don't daemonize, log to stdout.
      --disable-pin         don't check server public key against a known pin.
  -h, --help                show this help.
      --imaplib-debug       show raw IMAP messages (will include email password).
  -l, --log FILENAME        write log to FILENAME (default \"%(log)s\").
  -p, --pass=PASSFILE       use the email/password contained in PASSFILE. This file
                            should contain "[<imap_host>] <email> <password>" on a
                            single line, separated by whitespace. If <imap_host> is
                            omitted, it defaults to imap.(<email> domain):993.
      --pidfile FILENAME    write PID to FILENAME after daemonizing.
      --privdrop-user USER  switch UID and GID to those of USER.
      --unsafe-logging      don't scrub email password and IP addresses from logs.\
""" % {
    "progname": sys.argv[0],
    "log": DEFAULT_LOG_FILENAME,
}

def safe_str(s):
    """Return "[scrubbed]" if options.safe_logging is true, and s otherwise."""
    if options.safe_logging:
        return "[scrubbed]"
    else:
        return s

def log(msg):
    print >> options.log_file, (u"%s %s" % (time.strftime(LOG_DATE_FORMAT), msg)).encode("UTF-8")
    options.log_file.flush()

def main():
    opts, args = getopt.gnu_getopt(sys.argv[1:], "de:hi:l:p:", [
        "debug",
        "disable-pin",
        "email=",
        "help",
        "imap=",
        "imaplib-debug",
        "log=",
        "pass=",
        "pidfile=",
        "privdrop-user=",
        "unsafe-logging",
    ])
    for o, a in opts:
        if o == "-d" or o == "--debug":
            options.daemonize = False
            options.log_filename = None
        elif o == "--disable-pin":
            options.use_certificate_pin = False
        elif o == "-h" or o == "--help":
            usage()
            sys.exit()
        if o == "--imaplib-debug":
            options.imaplib_debug = True
        elif o == "-l" or o == "--log":
            options.log_filename = a
        elif o == "-p" or o == "--pass":
            options.password_filename = a
        elif o == "--pidfile":
            options.pid_filename = a
        elif o == "--privdrop-user":
            options.privdrop_username = a
        elif o == "--unsafe-logging":
            options.safe_logging = False

    if len(args) != 0:
        usage(sys.stderr)
        sys.exit(1)

    # Load the email password.
    if options.password_filename is None:
        print >> sys.stderr, "The --pass option is required."
        sys.exit(1)
    try:
        password_file = open(options.password_filename)
    except Exception, e:
        print >> sys.stderr, """\
    Failed to open password file "%s": %s.\
    """ % (options.password_filename, str(e))
        sys.exit(1)
    try:
        if not proc.check_perms(password_file.fileno()):
            print >> sys.stderr, "Refusing to run with group- or world-readable password file. Try"
            print >> sys.stderr, "\tchmod 600 %s" % options.password_filename
            sys.exit(1)
        for (lineno0, line) in enumerate(password_file.readlines()):
            line = line.strip("\n")
            if not line or line.startswith('#'): continue
            # we do this stricter regex match because passwords might have spaces in
            res = re.match(r"(?:(\S+)\s)?(\S+@\S+)\s(.+)", line)
            if not res:
                raise ValueError("could not find email or password on line %s" % (lineno0+1))
            (imap_addr_spec, email_addr, email_password) = res.groups()
            imap_addr = parse_addr_spec(
                imap_addr_spec or "", DEFAULT_IMAP_HOST, DEFAULT_IMAP_PORT)
            break
        else:
            raise ValueError("no email line found")
    except Exception, e:
        print >> sys.stderr, """\
    Failed to parse password file "%s": %s.
    Syntax is [<imap_host>] <email> <password>.
    """ % (options.password_filename, str(e))
        sys.exit(1)
    finally:
        password_file.close()

    if options.log_filename:
        options.log_file = open(options.log_filename, "a")
        # Send error tracebacks to the log.
        sys.stderr = options.log_file
    else:
        options.log_file = sys.stdout

    if options.daemonize:
        log(u"daemonizing")
        pid = os.fork()
        if pid != 0:
            if options.pid_filename:
                f = open(options.pid_filename, "w")
                print >> f, pid
                f.close()
            sys.exit(0)

    if options.privdrop_username is not None:
        log(u"dropping privileges to those of user %s" % options.privdrop_username)
        try:
            proc.drop_privs(options.privdrop_username)
        except BaseException, e:
            print >> sys.stderr, "Can't drop privileges:", str(e)
            sys.exit(1)

    if options.imaplib_debug:
        imaplib.Debug = 4

    login_limit = RateLimit()
    while True:
        try:
            imap = imap_login(imap_addr, email_addr, email_password)
            try:
                imap_loop(imap)
            except imaplib.IMAP4.error:
                imap.close()
                imap.logout()
        except (imaplib.IMAP4.error, ssl.SSLError, SSL.SSLError, socket.error), e:
            # Try again after a disconnection.
            log(u"lost server connection: %s" % str(e))
        except KeyboardInterrupt:
            break

        # Don't reconnect too fast.
        t = login_limit.time_to_wait()
        if t > 0:
            log(u"waiting %.2f seconds before logging in again" % t)
            time.sleep(t)

    log(u"closing")
    imap.close()
    imap.logout()

def message_get_date(msg):
    """Get the datetime when the message was received by reading the X-Received
    header, relative to UTC. Returns None on error."""
    x_received = msg["X-Received"]
    if x_received is None:
        log(u"X-Received is missing")
        return None
    try:
        _, date_str = x_received.rsplit(";", 1)
        date_str = date_str.strip()
    except ValueError:
        log(u"can't parse X-Received %s" % repr(x_received))
        return None
    date_tuple = email.utils.parsedate_tz(date_str)
    if date_tuple is None:
        log(u"can't parse X-Received date string %s" % repr(date_str))
        return None
    timestamp_utc = calendar.timegm(date_tuple[:8] + (0,)) - date_tuple[9]
    return datetime.datetime.utcfromtimestamp(timestamp_utc)

def message_ok(msg):
    date = message_get_date(msg)
    if date is not None:
        now = datetime.datetime.utcnow()
        age = time.mktime(now.utctimetuple()) - time.mktime(date.utctimetuple())
        if age > REGISTRATION_AGE_LIMIT:
            log(u"message dated %s UTC is too old: %d seconds" % (date, age))
            return False
        if -age > REGISTRATION_AGE_LIMIT:
            log(u"message dated %s UTC is from the future: %d seconds" % (date, -age))
            return False
    return True

def handle_message(msg):
    try:
        if fac.put_reg_proc(["fp-reg-decrypt"], msg.get_payload()):
            log(u"registered client")
        else:
            log(u"failed to register client")
    except Exception, e:
        log(u"error registering client")
        raise

def truncate_repr(s, n):
    if not isinstance(s, basestring):
        s = repr(s)
    if len(s) > n:
        return repr(s[:n]) + "[...]"
    else:
        return repr(s)
def check_imap_return(typ, data):
    if typ != "OK":
        raise imaplib.IMAP4.abort("Got type \"%s\": %s" % (typ, truncate_repr(data, 100)))

def imap_get_uid(imap, index):
    typ, data = imap.fetch(str(index), "(UID)")
    if data[0] is None:
        return None
    check_imap_return(typ, data)
    # This grepping for the UID is bogus, but imaplib doesn't properly break up
    # the list of name-value pairs for us.
    m = re.match(r'^\d+\s+\(.*\bUID\s+(\d+)\b.*\)\s*$', data[0])
    if m is None:
        raise imaplib.IMAP4.abort("Can't find UID in %s" % repr(data[0]))
    return m.group(1)

# Gmail's IMAP folders are funny: they are not real folders, but actually views
# of messages having a particular label. INBOX consists of messages having the
# INBOX label, for example. Deleting a message from a folder just removes its
# label, but the message itself continues to exist in "[Gmail]/All Mail".
#   https://support.google.com/mail/bin/answer.py?answer=78755
#   http://gmailblog.blogspot.com/2008/10/new-in-labs-advanced-imap-controls.html
# To really delete a message, you must copy it to "[Gmail]/Trash" and then
# delete it from there. Messages in Trash are deleted automatically after 30
# days, but we do it immediately.
def imap_loop(imap):
    while True:
        # Copy all messages to Trash, and work on them from there. This is a
        # failsafe so that messages will eventually be deleted if we are not
        # able to retrieve them. This act of copying also deletes from All Mail.
        typ, data = imap.select("[Gmail]/All Mail")
        check_imap_return(typ, data)
        imap.copy("1:*", "[Gmail]/Trash")

        typ, data = imap.select("[Gmail]/Trash")
        check_imap_return(typ, data)
        exists = int(data[0])
        if exists > 0:
            while True:
                # Grab message 1 on each iteration; remaining messages shift down so
                # the next message we process is also message 1.
                uid = imap_get_uid(imap, "1")
                if uid is None:
                    break

                typ, data = imap.uid("FETCH", uid, "(BODY[])")
                check_imap_return(typ, data)
                msg_text = data[0][1]
                typ, data = imap.uid("STORE", uid, "+FLAGS", "\\Deleted")
                check_imap_return(typ, data)
                typ, data = imap.expunge()
                check_imap_return(typ, data)

                try:
                    msg = email.message_from_string(msg_text)
                    if message_ok(msg):
                        handle_message(msg)
                except Exception, e:
                    log("Error processing message, deleting anyway: %s" % str(e))

        time.sleep(POLL_INTERVAL)

def imap_login(imap_addr, email_addr, email_password):
    """Make an IMAP connection, check the certificate and public key, and log in."""
    with keys.temp_cert(keys.PIN_GOOGLE_CA_CERT) as ca_certs_file:
        imap = IMAP4_SSL_REQUIRED(
            imap_addr[0], imap_addr[1], None, ca_certs_file.name)

    if options.use_certificate_pin:
        keys.check_certificate_pin(imap.ssl(), keys.PIN_GOOGLE_PUBKEY_SHA1)

    log(u"logging in as %s" % email_addr)
    imap.login(email_addr, email_password)

    return imap

class RateLimit(object):
    INITIAL_INTERVAL = 1.0
    # These constants are chosen to reach a steady state of one attempt every
    # ten minutes, assuming a new failing attempt after each penalty interval.
    MAX_INTERVAL = 10 * 60
    MULTIPLIER = 2.0
    DECAY = math.log(MULTIPLIER) / MAX_INTERVAL
    def __init__(self):
        self.time_last = time.time()
        self.interval = self.INITIAL_INTERVAL
    def time_to_wait(self):
        now = time.time()
        delta = now - self.time_last
        # Discount time already served.
        wait = max(self.interval - delta, 0)
        self.time_last = now
        self.interval = self.interval * math.exp(-self.DECAY * delta) * self.MULTIPLIER
        return wait

if __name__ == "__main__":
    main()
