1 files changed, 58 insertions, 20 deletions
diff --git a/tools/check-sth.py b/tools/check-sth.py
index 323d1a8..85e4661 100755
--- a/tools/check-sth.py
+++ b/tools/check-sth.py
@@ -7,6 +7,8 @@
 import sys
 import argparse
 import json
+import errno
+import shutil
 from datetime import datetime, timedelta
 from certtools import get_sth
 
@@ -15,30 +17,47 @@ NAGIOS_WARN = 1
 NAGIOS_CRIT = 2
 NAGIOS_UNKNOWN = 3
 
+DEFAULT_CUR_FILE = 'cur-sth.json'
+DEFAULT_PREV_FILE = 'prev-sth.json'
+
 parser = argparse.ArgumentParser(description="")
-parser.add_argument('--sth-file', help="File with previous STH to verify against")
+parser.add_argument('--cur-sth',
+                    default=DEFAULT_CUR_FILE,
+                    help="File containing current STH (default=%s)" % DEFAULT_CUR_FILE)
+parser.add_argument('--prev-sth',
+                    default=DEFAULT_PREV_FILE,
+                    help="File containing previous STH (default=%s" % DEFAULT_PREV_FILE)
 parser.add_argument('baseurl', help="Base URL for CT server")
-args = parser.parse_args()
 
 def print_sth(sth):
-    print sth['timestamp']
-    print sth['sha256_root_hash']
-    print sth['tree_size']
-    print sth['tree_head_signature']
+    if sth is None:
+        print "NONE"
+    else:
+        print sth['timestamp']
+        print sth['sha256_root_hash']
+        print sth['tree_size']
+        print sth['tree_head_signature']
 
-def get_cur_sth(baseurl):
+def get_new_sth(baseurl):
     try:
-        sth = get_sth(args.baseurl)
+        sth = get_sth(baseurl)
     except Exception, e:
         print e
         sys.exit(NAGIOS_UNKNOWN)
     return sth
         
-def read_prev_sth(filename):
-    if args.sth_file is not None:
-        prev_sth = json.loads(open(args.sth_file).read())
-        return prev_sth
-    return None
+def read_sth(fn):
+    try:
+        f = open(fn)
+    except IOError, errno.ENOENT:
+        return None
+    return json.loads(f.read())
+
+def mv_file(fromfn, tofn):
+    shutil.move(fromfn, tofn)
+
+def write_file(fn, sth):
+    open(fn, 'w').write(json.dumps(sth))
 
 def check_age(sth):
     now = datetime.now()
@@ -57,11 +76,30 @@ def check_treesize(cur, prev):
               (cur['tree_size'], prev['tree_size'])
             sys.exit(NAGIOS_CRIT)
 
-cur_sth = get_cur_sth(args.baseurl)
-prev_sth = read_prev_sth(args.sth_file)
-check_age(cur_sth)
-check_treesize(cur_sth, prev_sth)
-# TODO: verify signature
-# TODO: get and verify consistency proof
+def main(args):
+    if args.cur_sth is None:
+        args.cur_sth = "cur-sth.json"
+    if args.prev_sth is None:
+        args.prev_sth = "prev-sth.json"
+
+    new = get_new_sth(args.baseurl)
+    cur = read_sth(args.cur_sth)
+    if cur is None or new['sha256_root_hash'] != cur['sha256_root_hash']:
+        if cur is not None:
+            mv_file(args.cur_sth, args.prev_sth)
+        write_file(args.cur_sth, new)
+        cur = new
+    prev = read_sth(args.prev_sth)
+
+    #print_sth(cur)
+    #print_sth(prev)
+
+    check_age(cur)
+    check_treesize(cur, prev)
+    # TODO: verify signature
+    # TODO: get and verify consistency proof
+
+    sys.exit(NAGIOS_OK)
 
-sys.exit(NAGIOS_OK)
+if __name__ == '__main__':
+    main(parser.parse_args())