Handle db transactions consistently

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)