#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
This script doesn't actually do much that's useful.

It was designed to be an orchestration test between a few libraries,
now quite a few libraries.
"""

import os
import time
import sys
import getopt
import logging
import pickle
import random
import datetime
import uuid
logging.basicConfig(level=logging.DEBUG)

from pydisque.client import Client
from pybloomd import BloomdClient
from pyhlld import HlldClient

import redis

from itsdangerous import Serializer

from Crypto.PublicKey import RSA
from Crypto.Signature import PKCS1_v1_5
from Crypto.Hash import SHA
from Crypto import Random

import leveldb

def usage():
    """Print the usage of the software."""
    txt = """
    objbomber.py [args]

    acceptable arguments:
    --mode=["generate", "listen", "check"]
        generate - tell the script to generate some random messages
        listen - tell the script to listen to disque for messages
        check - tell the script to validate messages against bloomd
        initialize - tell the script that this is the inital run
                     (sets up cryptographic keys)
    --dservers=[comma separated host:port list of disque servers]
        defaults to: "localhost:7712,localhost:7711"
        (can be a single host:port pair)
    --dqueue=[name of queue to use in Disque | default "objbomber"]
        (note that "{dqueue}.check" will also be used)
    --secret=[the shared secret | default "coolsecretbro"]
    --bserver=[single host:port:udpport of bloomd the server]
        defaults to: localhost:8673
    --hserver=[single host:port]
        defaults to localhost:4553
    --rserver=[single host:port]
        defaults to localhost:6379
    --bfiltername=[single name of the bloom filter | default "objbomber"]
    --sleep=[seconds to wait in our core loop | default 0]
    --help (prints this message)
    """
    print(txt)


def main():
    """Primary CLI application logic."""
    try:
        opts, args = getopt.getopt(sys.argv[1:], "h:v", ["help",
                                                         "dservers=",
                                                         "dqueue=",
                                                         "secret=",
                                                         "bserver=",
                                                         "hserver=",
                                                         "rserver=",
                                                         "rchannel=",
                                                         "bfiltername=",
                                                         "hllname=",
                                                         "mode=",
                                                         "sleep="])
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit()

    modes = ("generate", "listen", "check", "adaptive", "initialize", "subscriber")

    # set defaults
    mode = None
    dservers = "localhost:7712,localhost:7711"
    dqueue = "objbomber"
    secret = "coolsecretbro"
    bserver = "localhost:8673"
    hserver = "localhost:4553"
    rserver = "localhost:6379"
    rchannel = "objbomber"
    bfiltername = "objbomber"
    hllname = "objbomber"
    sleep = None
    userhomedir = os.path.expanduser("~")

    # flippin' switches...
    for o, a in opts:
        if o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("--dservers"):
            dservers = a
        elif o in ("--queue"):
            dqueue = [a]
        elif o in ("--secret"):
            secret = a
        elif o in ("--bserver"):
            bserver = a
        elif o in ("--hserver"):
            hserver = a
        elif o in ("--rserver"):
            rserver = a
        elif o in ("--rchannel"):
            rchannel = a
        elif o in ("--bfiltername"):
            bfiltername = a
        elif o in ("--hllname"):
            hllname = a
        elif o in ("--mode"):
            if a in modes:
                mode = a
            else:
                usage()
                sys.exit()
        elif o in ("--listen"):
            mode_listen = True
        elif o in ("--check"):
            mode_check = True
        elif o in ("--sleep"):
            sleep = int(a)
        else:
            assert False, "unhandled option"

    checkdqueue = dqueue + ".check"

    if sleep in (None, 0):
        sleep = 0.0001

    # mode must be set
    if not mode:
        usage()
        sys.exit()

    # Handler for the cryptographic signatures
    # TODO: secret should be "secret" + a version number
    s = Serializer(secret)

    # config basics
    datadir = userhomedir + "/.objbomber"

    # prepare servers and queue lists
    dservers = dservers.split(",")
    bserver = [bserver]
    hserver = hserver

    # all modes use Disque
    logging.info("Connecting to Disque...")
    disque_client = Client(dservers)
    disque_client.connect()

    if mode in ("check", "listen"):
        logging.info("Creating Bloomd Client...")
        bloomd_client = BloomdClient(bserver)
        bfilter = bloomd_client.create_filter(bfiltername)

        # add pyhlld
        logging.info("Creating HLLD Client... - not yet used")
        hlld_client = HlldClient(hserver)
        hll = hlld_client.create_set(hllname)

    if mode in ("check", "listen", "generate", "subscriber"):
        # add redis hll & pubsub
        logging.info("Creating Redis Client...")
        rhost, rport = rserver.split(":")
        redd = redis.StrictRedis(host=rhost, port=rport, db=0)
        redpubsub = redd.pubsub()

    if mode in ("subscriber"):
        redpubsub.subscribe(rchannel)

    if mode in ("generate"):
        # TODO: check on how well LevelDB handles
        #       multiple clients
        db = leveldb.LevelDB(datadir + '/notary')

    # special mode to handle our first run
    # TODO: push into a function
    # TODO: handle filesystem errors
    # TODO: reconsider using Cement for all of this
    # TODO: generate an instance GUID
    if mode == "initialize":

        UUID = uuid.uuid4()
        logging.info("Our system UUID is now: %s" % UUID)
        # TODO: save and load this uuid

        # check to see if there is a ~/.objbomber directory, quit if there is
        # TODO: this does not handle errors in initalization
        logging.info("Checking for .objbomber in %s..." % userhomedir)
        if os.path.exists(datadir):
            logging.info("Already been initialized!")
            # TODO: print some information about how to handle this
            sys.exit()

        # TODO: make one
        os.mkdir(datadir, 0700)

        # generate our RSA signing key
        # TODO: make # of bits in key a flag
        logging.info("Begining to create our notary key.")
        logging.info("Reading from RNG.")
        rng = Random.new().read

        logging.info("Generating RSA key...")
        privRSAkey = RSA.generate(4096, rng)
        privPEM = privRSAkey.exportKey()
        pubRSAkey = privRSAkey.publickey()
        pubPEM = pubRSAkey.exportKey()
        logging.info("Key generated.")

        # save privkey to disk
        with open(datadir + "/privkey.pem", "w") as keyfile:
            keyfile.write(privPEM)
            keyfile.close()
            os.chmod(datadir + "/privkey.pem", 0700)
            logging.info("Unencrypted RSA key written to disk.") 

        # save the pubkey
        with open(datadir + "/pubkey.pem", "w") as keyfile:
            keyfile.write(pubPEM)
            keyfile.close()
            logging.info("Public RSA key written to disk.")

        logging.info("Creating crypto notary storage.")
        leveldb.LevelDB(datadir + '/notary')

        # we absolutely must quit here, or we will get stuck in
        # an infinate loop
        sys.exit()

    # load our secret key (TODO: this is probably better as try/exc)
    # and build our contexts
    with open(datadir + "/privkey.pem", "r") as keyfile:
        privRSAkey = RSA.importKey(keyfile.read())

    while True:
        # TODO: Adaptive Mode - this mode should peek the queues, and
        # make a decision about where this thread can make the most
        # impact on its own.
        if mode == "adaptive":

            # TODO: Do some queue peeking.

            # TODO: Make some decisions about which mode to adapt to.

            pass

        # TODO: All modes should be placed into functions.

        # Listen Mode - Listens to the queue, pulls out jobs,
        #   validates the signature, puts them in bloomd
        if mode == "listen":
            logging.info("Getting Jobs from Disque.")
            jobs = disque_client.get_job([dqueue])

            print("Got %d jobs." % len(jobs))

            for queue_name, job_id, job in jobs:
                logging.debug("Handling a job: %s" % job)

                try:
                    job = s.loads(job)
                    logging.debug("Job Authenticated: %s" % job)
                except:
                    logging.warning("Job did not pass authentication.")
                    disque_client.nack_job(job_id)

                # add to bloom filter
                try:
                    bfilter.add(job)
                except:
                    logging.warning("Job was not added to bloomd.")
                    disque_client.nack_job(job_id)

                try:
                    hllResponse = hll.add(job)
                except:
                    logging.warning("Job was not added to hlld.")
                    disque_client.nack_job(job_id)

                # TODO: add redis HLL support

                # tell disque that this job has been processed
                disque_client.ack_job(job_id)

                # sign the check job
                job = s.dumps(job)

                # throw this message on the check queue
                disque_client.add_job(checkdqueue, job)

        elif mode == "check":
            # TODO
            # Check the secondary disque queue for checks
            # Ask the bloom filter if they have seen this

            logging.info("Getting Jobs from Disque.")
            jobs = disque_client.get_job([checkdqueue])

            for queue_name, job_id, job in jobs:
                logging.debug("Checking: %s" % job)

                try:
                    job = s.loads(job)
                except:
                    disque_client.nack_job(job_id)

                # we don't NACK on failed cache hits

                if job in bfilter:
                    logging.info("Confirming: %s" % job)
                else:
                    logging.info("Not found in bloom filter: %s" % job)

                disque_client.ack_job(job_id)

        elif mode == "generate":
            # TODO - where will these messages come from?
            # for now they will just be random numbers, but
            # really we should make them objects to really be
            # testing serialization
            msg = [random.randint(1000, 1000000),
                   random.randint(1000, 1000000)]

            # itsdangerous serialization & signing
            msg = s.dumps(msg)

            # Now that this message is serialized, we can sign it again with a
            # public key.
            # TODO: incorporate saving the signature into the notary records
            msghash = SHA.new(msg)
            signer = PKCS1_v1_5.new(privRSAkey)
            signature = signer.sign(msghash)
            assert signer.verify(msghash, signature)

            record = {'message': msg, 'signature': signature}
            record = pickle.dumps(record)

            # send the job over to Disque
            # TODO: add more command line flags for queuing
            job_id = disque_client.add_job(dqueue, msg)
            logging.debug("Added a job to Disque: %s" % msg)

            # publish just the signature on redis pubsub
            redd.publish(rchannel, signature)

            # TODO: save the publication in the notary
            # TODO: do more then just save the signatures
            # TODO: add a GUID to the key
            key = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f")
            db.Put(key, record)

            # testing the results of leveldb's store, this is a test, and
            # an expensive test
            sig2 = db.Get(key)
            sig2 = pickle.loads(sig2)['signature']
            assert signer.verify(msghash, sig2)

        elif mode == "subscriber":
            msg = redpubsub.get_message()
            # TODO: do something useful, like log
            if msg:
                print("got a message")

        time.sleep(sleep)

if __name__ == "__main__":
    main()
