Permalink
Browse files

Handle db transactions consistently

  • Loading branch information...
jgor committed Jan 15, 2019
1 parent 771bed7 commit 6a87a19dc3a860f76563749fb294449c568e63bc
Showing with 24 additions and 37 deletions.
  1. +24 −37 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)

0 comments on commit 6a87a19

Please sign in to comment.