From 6a87a19dc3a860f76563749fb294449c568e63bc Mon Sep 17 00:00:00 2001 From: jgor Date: Tue, 15 Jan 2019 11:45:05 -0600 Subject: [PATCH] Handle db transactions consistently --- dorkbot.py | 61 +++++++++++++++++++++--------------------------------- 1 file changed, 24 insertions(+), 37 deletions(-) diff --git a/dorkbot.py b/dorkbot.py index 0e6ebd0..d908855 100755 --- a/dorkbot.py +++ b/dorkbot.py @@ -9,6 +9,7 @@ from urllib.parse import urlparse except ImportError: from urlparse import urlparse +from contextlib import closing import datetime import hashlib import importlib @@ -194,21 +195,18 @@ def __init__(self, database): sys.exit(1) try: - c = self.db.cursor() - c.execute("CREATE TABLE IF NOT EXISTS targets (url VARCHAR PRIMARY KEY)") - c.execute("CREATE TABLE IF NOT EXISTS fingerprints (fingerprint VARCHAR PRIMARY KEY)") - self.db.commit() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("CREATE TABLE IF NOT EXISTS targets (url VARCHAR PRIMARY KEY)") + c.execute("CREATE TABLE IF NOT EXISTS fingerprints (fingerprint VARCHAR PRIMARY KEY)") except self.module.Error as e: print("ERROR loading database - %s" % e, file=sys.stderr) sys.exit(1) def get_targets(self): try: - c = self.db.cursor() - c.execute("SELECT url FROM targets") - targets = [row[0] for row in c.fetchall()] - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("SELECT url FROM targets") + targets = [row[0] for row in c.fetchall()] except self.module.Error as e: print("ERROR getting targets - %s" % e, file=sys.stderr) sys.exit(1) @@ -217,10 +215,9 @@ def get_targets(self): def get_next_target(self): try: - c = self.db.cursor() - c.execute("SELECT url FROM targets LIMIT 1") - row = c.fetchone() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("SELECT url FROM targets LIMIT 1") + row = c.fetchone() except self.module.Error as e: print("ERROR getting next target - %s" % e, file=sys.stderr) sys.exit(1) @@ -230,10 +227,9 @@ def get_next_target(self): def get_random_target(self): try: - c = self.db.cursor() - c.execute("SELECT url FROM targets ORDER BY RANDOM() LIMIT 1") - row = c.fetchone() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("SELECT url FROM targets ORDER BY RANDOM() LIMIT 1") + row = c.fetchone() except self.module.Error as e: print("ERROR getting random target - %s" % e, file=sys.stderr) sys.exit(1) @@ -243,10 +239,8 @@ def get_random_target(self): def add_target(self, url): try: - c = self.db.cursor() - c.execute("%s INTO targets VALUES (%s)" % (self.insert, self.param), (url,)) - self.db.commit() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("%s INTO targets VALUES (%s)" % (self.insert, self.param), (url,)) except self.module.IntegrityError as e: if "UNIQUE constraint failed" in e: return @@ -257,20 +251,17 @@ def add_target(self, url): def delete_target(self, url): try: - c = self.db.cursor() - c.execute("DELETE FROM targets WHERE url=(%s)" % self.param, (url,)) - self.db.commit() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("DELETE FROM targets WHERE url=(%s)" % self.param, (url,)) except self.module.Error as e: print("ERROR deleting target - %s" % e, file=sys.stderr) sys.exit(1) def get_scanned(self, fingerprint): try: - c = self.db.cursor() - c.execute("SELECT fingerprint FROM fingerprints WHERE fingerprint = (%s)" % self.param, (fingerprint,)) - row = c.fetchone() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("SELECT fingerprint FROM fingerprints WHERE fingerprint = (%s)" % self.param, (fingerprint,)) + row = c.fetchone() except self.module.Error as e: print("ERROR looking up fingerprint - %s" % e, file=sys.stderr) sys.exit(1) @@ -280,20 +271,16 @@ def get_scanned(self, fingerprint): def add_fingerprint(self, fingerprint): try: - c = self.db.cursor() - c.execute("%s INTO fingerprints VALUES (%s)" % (self.insert, self.param), (fingerprint,)) - self.db.commit() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("%s INTO fingerprints VALUES (%s)" % (self.insert, self.param), (fingerprint,)) except self.module.Error as e: print("ERROR adding fingerprint - %s" % e, file=sys.stderr) sys.exit(1) def flush_fingerprints(self): try: - c = self.db.cursor() - c.execute("DELETE FROM fingerprints") - self.db.commit() - c.close() + with self.db, closing(self.db.cursor()) as c: + c.execute("DELETE FROM fingerprints") except self.module.Error as e: print("ERROR flushing fingerprints - %s" % e, file=sys.stderr) sys.exit(1)