using two databases

4 files changed, 76 insertions, 47 deletions

diff --git a/monitor/josef_experimental.py b/monitor/josef_experimental.py
index 433bb77..c9605bd 100755
--- a/monitor/josef_experimental.py
+++ b/monitor/josef_experimental.py
@@ -27,15 +27,16 @@ monitored_domains = [
 ]
 
 
-db = db_open("./tmpdatabase")
+# db = db_open("./tmpdatabase")
 
 # db_add_domain_2(db, "www.cox.a.com", "{dummydata}")
 # print db_lookup_domain(db, "www.cox.a.com")
 # db_add_domain_2(db, "www.cox.b.com", "{dummydata}")
 # print db_lookup_domain(db, "www.cox.b.com")
 # print db.Get("www.cox.b.com")
-db_add_domain_2(db, "www.cox.a.com", "{dummydata3}")
-print db_lookup_domain(db, "www.cox.a.com")
+db_add_domain_2("./tmpdatabase/", "www.cox.a.com", "{dummydata3}")
+db_add_domain_2("./tmpdatabase/", "www.ct.nordu.net", "{dummydata1}")
+print db_lookup_domain("./tmpdatabase/", "www.cox.a.com")
 # print db.Get("com")
 # print db.Get("a.com")
 # print db.Get("cox.a.com")
diff --git a/monitor/josef_leveldb.py b/monitor/josef_leveldb.py
index 5c7fc8a..b4b3255 100755
--- a/monitor/josef_leveldb.py
+++ b/monitor/josef_leveldb.py
@@ -8,6 +8,8 @@ import leveldb
 SEP = ";"
 # db = None
 
+dbs = {}
+
 
 def match_domain(d1, d2):
     # Exact match
@@ -31,8 +33,14 @@ def match_domain(d1, d2):
 
 
 def db_open(fn='./cert_db'):
-    db = leveldb.LevelDB(fn)
-    return db
+    global dbs
+    if fn in dbs:
+        return dbs[fn]
+    else:
+        print "Opening " + fn
+        db = leveldb.LevelDB(fn)
+        dbs[fn] = db
+        return db
 
 def db_append(db, key, val):
     if db is None:
@@ -50,11 +58,25 @@ def db_append(db, key, val):
         tmpl.append(val)
         db.Put(key,SEP.join(tmpl))
 
-def db_add_domain_2(db, key, val):
-    if db is None:
-        print "ERROR: NO DATABASE SET!"
-        return
+def db_add_domain_2(db_dir, key_in, val, db_in=None):
+    if db_in == None:
+        key_list  = key_in.split('/')[0].split(',')[0].split('.')
+        if key_list[-1] == "com":
+            key  = '.'.join(key_list[:-1])
+            db = db_open(db_dir + key_list[-1])
+        else:
+            key  = '.'.join(key_list)
+            db = db_open(db_dir + "other")
+        # key  = '.'.join(key_list)
+        # db = db_open(db_dir + key_list[-1])
+        # key  = '.'.join(key_list)
+        # db = db_open(db_dir)
+        # print "DB: " + db_dir + key_list[-1]
+    else:
+        db = db_in 
+        key = key_in
 
+    # print "key: ", key, " val: ", val
     try:
         tmp = db.Get(key)
         tmpl = tmp.split(SEP)
@@ -67,7 +89,7 @@ def db_add_domain_2(db, key, val):
         db.Put(key,SEP.join([val]))
         new_key = key.split('.',1)[-1]
         if new_key != key:
-            db_add_domain_2(db,new_key,key)
+            db_add_domain_2(db,new_key,key, db)
 
 
 def db_add_domain(db, domain, data):
@@ -87,32 +109,33 @@ def db_add_domain(db, domain, data):
     db_append(db, k, data)
 
 
-def db_add_certs(db, data):
-    if db is None:
+def db_add_certs(db_dir, data):
+    if db_dir is None:
         print "ERROR: NO DATABASE SET!"
         return
 
-    # batch = leveldb.WriteBatch()
     for cert in data:
         try:
-            db_add_domain_2(db, cert["subject"].split("CN=")[1], str(cert))
+            db_add_domain_2(db_dir, cert["subject"].split("CN=")[1], str(cert))
         except IndexError:
             pass
         try:
             for line in cert["SAN"].split("DNS:")[1:]:
-                db_add_domain_2(db, line, str(cert))
+                db_add_domain_2(db_dir, line, str(cert))
         except KeyError:
             pass
         except IndexError:
             pass
-    # db.Write(batch, sync = True)
 
 
-def db_lookup_domain(db, domain):
+def db_lookup_domain(db_dir, domain):
     domain_list = domain.split('.')
     res = []
+    db = db_open(db_dir + domain_list[-1])
+    domain_list = domain_list[:-1]
 
     cur_domain = domain_list.pop()
+    # print cur_domain
     intermediate = db.Get(cur_domain).split(SEP)
     last = False
     one_more = False
diff --git a/monitor/josef_monitor.py b/monitor/josef_monitor.py
index 3a71987..14a280a 100755
--- a/monitor/josef_monitor.py
+++ b/monitor/josef_monitor.py
@@ -26,7 +26,10 @@ else:
 if not os.path.exists(OUTPUT_DIR):
     os.makedirs(OUTPUT_DIR)
 
-DB = None
+if not os.path.exists(DB_PATH):
+    os.makedirs(DB_PATH)
+
+# DB = None
 
 parser = argparse.ArgumentParser(description="")
 
@@ -51,7 +54,7 @@ class ctlog:
         start_size = self.entries
         while self.entries < self.sth["tree_size"]:
             tmp_size = self.entries
-            self.subtree, self.entries = fetch_and_increment_subtree(self.entries, self.sth["tree_size"] -1, self.url, self.subtree)
+            self.subtree, self.entries = self.fetch_and_increment_subtree(self.entries, self.sth["tree_size"] -1, self.url, self.subtree)
             if tmp_size != self.entries:
                 self.log("Got entries " + str(tmp_size) + " to " \
                 + str(self.entries -1 ) + " of " + str(self.sth["tree_size"]-1))
@@ -66,6 +69,29 @@ class ctlog:
         #     self.log("No new entries.")
 
 
+    def fetch_and_increment_subtree(self, first, last, url, subtree =[[]]):
+        # global DB
+        # try:
+            new_leafs = []
+            if first <= last:
+                entries = get_entries(url, first, last)["entries"]
+                tmp_cert_data = []
+                for item in entries:
+                    tmp_data = check_domain(item, url)
+                    if tmp_data:
+                        tmp_cert_data.append(tmp_data)
+                    new_leafs.append(get_leaf_hash(base64.b64decode(item["leaf_input"])))
+                if DB_PATH:
+                    self.log("Adding to database...")
+                    db_add_certs(DB_PATH, tmp_cert_data)
+                    self.log("done adding to DB.")
+                if DEFAULT_CERT_FILE:
+                    append_file(DEFAULT_CERT_FILE, tmp_cert_data)
+                subtree = reduce_tree(new_leafs, subtree)
+        # except:
+        #     print "Failed to build subtree :("
+            return subtree, len(new_leafs) + first
+
 
     def to_dict(self):
         d = {}
@@ -215,28 +241,6 @@ def check_domain(raw_entry, log=None):
     except IndexError:
         return None
 
-def fetch_and_increment_subtree(first, last, url, subtree =[[]]):
-        global DB
-    # try:
-        new_leafs = []
-        if first <= last:
-            entries = get_entries(url, first, last)["entries"]
-            tmp_cert_data = []
-            for item in entries:
-                tmp_data = check_domain(item, url)
-                if tmp_data:
-                    tmp_cert_data.append(tmp_data)
-                new_leafs.append(get_leaf_hash(base64.b64decode(item["leaf_input"])))
-            if DEFAULT_CERT_FILE:
-                if DB is None:
-                    append_file(DEFAULT_CERT_FILE, tmp_cert_data)
-                else:
-                    db_add_certs(DB, tmp_cert_data)
-            subtree = reduce_tree(new_leafs, subtree)
-    # except:
-    #     print "Failed to build subtree :("
-        return subtree, len(new_leafs) + first
-
 def verify_subtree(sth, subtree, base_url):
     try:
         tmp = deepcopy(subtree)
@@ -318,13 +322,12 @@ def get_all_roots(base_url):
             root_cert = base64.decodestring(accepted_cert)
             print get_cert_info(root_cert)["subject"]
 
-
 def main(args):
     global DB
     # Create logs
     logs = []
-    if DB_PATH:
-        DB = db_open(DB_PATH)
+    # if DB_PATH:
+    #     DB = db_open(DB_PATH)
     try:
         for item in ctlogs:
             logs.append(ctlog(item, ctlogs[item][0], ctlogs[item][1]))
diff --git a/monitor/monitor_conf.py b/monitor/monitor_conf.py
index 69fc5ba..1f51a10 100644
--- a/monitor/monitor_conf.py
+++ b/monitor/monitor_conf.py
@@ -7,10 +7,12 @@ INTERVAL = 30
 OUTPUT_DIR = "output/"
 
 # Output file for certificate data.
-# Set to None to disable
-DEFAULT_CERT_FILE = OUTPUT_DIR + "cert_data.json"
+# Set to None to disable textfile writing
+DEFAULT_CERT_FILE = None
+# DEFAULT_CERT_FILE = OUTPUT_DIR + "cert_data.json"
 
-DB_PATH = './cert_db'
+# Set to None to disable database output
+DB_PATH = './tmpdb/'
 
 # CT logs and associated keys
 ctlogs = {