class TestHlld(unittest.TestCase):
"""TestCase class for pyhlld."""
testID = None
def setUp(self):
"""Setup the tests."""
# Create a client to a local hlld server, default port
self.client = HlldClient('localhost')
self.testID = "%d.%d" % (time.time(),
random.randint(1000, 1000000))
def test_simple(self):
"""Test the most basic functionality of pyhlld."""
# Get or create the foobar set
foobar = self.client.create_set("foobar")
# Set a key and check the size
foobar.add("Test Key!")
assert int(foobar.info()["size"]) == 1
def test_pipeline(self):
# Get or create the pipe set
pipe = self.client.create_set("pipe").pipeline()
# Chain multiple add commands
results = pipe.add("foo").add("bar").add("baz").execute()
assert results[0]
assert results[1]
assert results[2]
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)
def setUp(self):
"""Setup the tests."""
# Create a client to a local hlld server, default port
self.client = HlldClient('localhost')
self.testID = "%d.%d" % (time.time(),
random.randint(1000, 1000000))
