Postgresql binary data - properly insert and read

pwman/data/database.py

@@ -122,7 +122,6 @@ class Database(object):
     def _get_tag(self, tagcipher):
         sql_search = "SELECT * FROM TAG"
         self._cur.execute(sql_search)
-
         ce = CryptoEngine.get()
 
         try:
@@ -144,7 +143,7 @@ class Database(object):
         if rv:
             return rv
         else:
-            self._cur.execute(self._insert_tag_sql, ([tagcipher]))
+            self._cur.execute(self._insert_tag_sql, list(map(self._data_wrapper, (tagcipher,))))  # noqa
             try:
                 return self._cur.fetchone()[0]
             except TypeError:
@@ -197,7 +196,7 @@ class Database(object):
     def add_node(self, node):
         node_tags = list(node)
         node, tags = node_tags[:4], node_tags[-1]
-        self._cur.execute(self._add_node_sql, (node))
+        self._cur.execute(self._add_node_sql, list(map(self._data_wrapper, (node))))  # noqa
         try:
             nid = self._cur.fetchone()[0]
         except TypeError:
@@ -252,7 +251,7 @@ class Database(object):
         self._cur.execute("INSERT INTO CRYPTO VALUES({}, {})".format(self._sub,
                                                                      self._sub),  # noqa
 
-                          (seed, digest))
+                          list(map(self._data_wrapper, (seed, digest))))
         self._con.commit()
 
     def loadkey(self):
@@ -263,7 +262,7 @@ class Database(object):
         try:
             self._cur.execute(sql)
             seed, digest = self._cur.fetchone()
-            return seed + u'$6$' + digest
+            return seed + '$6$' + digest
         except TypeError:  # pragma: no cover
             return None
 
@@ -272,9 +271,9 @@ class Database(object):
         sql = "INSERT INTO CRYPTO(SEED, DIGEST) VALUES({},{})".format(self._sub,  # noqa
                                                                       self._sub)  # noqa
         self._cur.execute("DELETE FROM CRYPTO")
-        self._cur.execute(sql, (salt, digest))
-        self._digest = digest.encode('utf-8')
-        self._salt = salt.encode('utf-8')
+        self._cur.execute(sql, list(map(self._data_wrapper, (salt, digest))))
+        self._digest = digest.encode()
+        self._salt = salt.encode()
         self._con.commit()
 
     def close(self):  # pragma: no cover

pwman/data/drivers/mysql.py

@@ -62,6 +62,7 @@ class MySQLDatabase(Database):
                               "NOTES) "
                               "VALUES(%s, %s, %s, %s)")
         self._insert_tag_sql = "INSERT INTO TAG(DATA) VALUES(%s)"
+        self._data_wrapper = lambda x: x
         self.ProgrammingError = mysql.ProgrammingError
 
     def _open(self):

pwman/data/drivers/postgresql.py

@@ -21,7 +21,9 @@
 
 """Postgresql Database implementation."""
 import psycopg2 as pg
+
 from pwman.data.database import Database, __DB_FORMAT__
+from pwman.util.crypto_engine import CryptoEngine
 
 
 class PostgresqlDatabase(Database):
@@ -68,9 +70,117 @@ class PostgresqlDatabase(Database):
                               'NOTES) VALUES(%s, %s, %s, %s) RETURNING ID')
         self._insert_tag_sql = "INSERT INTO TAG(DATA) VALUES(%s) RETURNING ID"
         self.ProgrammingError = pg.ProgrammingError
+        self._data_wrapper = lambda x: pg.Binary(x)
 
     def _open(self):
 
         self._con = pg.connect(self._pgsqluri.geturl())
         self._cur = self._con.cursor()
         self._create_tables()
+
+    def _get_tag(self, tagcipher):
+        sql_search = "SELECT * FROM TAG"
+        self._cur.execute(sql_search)
+        ce = CryptoEngine.get()
+
+        try:
+            tag = ce.decrypt(tagcipher)
+            encrypted = True
+        except Exception:
+            tag = tagcipher
+            encrypted = False
+
+        rv = self._cur.fetchall()
+        for idx, cipher in rv:
+            cipher = cipher.tobytes()
+            if encrypted and tag == ce.decrypt(cipher):
+                return idx
+            elif tag == cipher:
+                return idx
+
+    def _create_tables(self):
+        if self._check_tables():
+            return
+        try:
+            self._cur.execute("CREATE TABLE NODE(ID SERIAL PRIMARY KEY, "
+                              "USERNAME BYTEA NOT NULL, "
+                              "PASSWORD BYTEA NOT NULL, "
+                              "URL BYTEA NOT NULL, "
+                              "NOTES BYTEA NOT NULL"
+                              ")")
+
+            self._cur.execute("CREATE TABLE TAG"
+                              "(ID  SERIAL PRIMARY KEY,"
+                              "DATA BYTEA NOT NULL)")
+
+            self._cur.execute("CREATE TABLE LOOKUP ("
+                              "nodeid INTEGER NOT NULL REFERENCES NODE(ID),"
+                              "tagid INTEGER NOT NULL REFERENCES TAG(ID)"
+                              ")")
+
+            self._cur.execute("CREATE TABLE CRYPTO "
+                              "(SEED BYTEA, DIGEST BYTEA)")
+
+            self._cur.execute("CREATE TABLE DBVERSION("
+                              "VERSION TEXT NOT NULL)")
+
+            self._cur.execute("INSERT INTO DBVERSION VALUES(%s)",
+                              (self.dbversion,))
+
+            self._con.commit()
+        except Exception:  # pragma: no cover
+            self._con.rollback()
+
+    def savekey(self, key):
+        salt, digest = key.split('$6$')
+        try:
+            salt, digest = salt.encode(), digest.encode()
+        except AttributeError:
+            pass
+
+        sql = "INSERT INTO CRYPTO(SEED, DIGEST) VALUES({},{})".format(self._sub,  # noqa
+                                                                      self._sub)  # noqa
+        self._cur.execute("DELETE FROM CRYPTO")
+        self._cur.execute(sql, list(map(self._data_wrapper, (salt, digest))))
+        self._digest = digest
+        self._salt = salt
+        self._con.commit()
+
+    def loadkey(self):
+        """
+        return _keycrypted
+        """
+        sql = "SELECT * FROM CRYPTO"
+        try:
+            self._cur.execute(sql)
+            seed, digest = self._cur.fetchone()
+            return seed.tobytes() + b'$6$' + digest.tobytes()
+        except TypeError:  # pragma: no cover
+            return None
+
+    def getnodes(self, ids):
+        if ids:
+            sql = ("SELECT * FROM NODE WHERE ID IN ({})"
+                   "".format(','.join(self._sub for i in ids)))
+        else:
+            sql = "SELECT * FROM NODE"
+        self._cur.execute(sql, (ids))
+        nodes = self._cur.fetchall()
+        if not nodes:
+            return []
+
+        nodes_w_tags = []
+        for node in nodes:
+            tags = [t.tobytes() for t in self._get_node_tags(node)]
+            nodes_w_tags.append([node[0]] + [item.tobytes() for item in node[1:]] + tags)
+
+        return nodes_w_tags
+
+    def listtags(self):
+        self._clean_orphans()
+        get_tags = "select DATA from TAG"
+        self._cur.execute(get_tags)
+        tags = self._cur.fetchall()
+        if tags:
+            return [t[0].tobytes() for t in tags]
+        return []  # pragma: no cover

pwman/data/drivers/sqlite.py

@@ -54,6 +54,7 @@ class SQLite(Database):
         self._list_nodes_sql = "SELECT NODEID FROM LOOKUP WHERE TAGID = ? "
         self._insert_tag_sql = "INSERT INTO TAG(DATA) VALUES(?)"
         self._sub = '?'
+        self._data_wrapper = lambda x: x
 
     def _open(self):
         self._con = sqlite.connect(self._filename)

tests/test_postgresql.py

@@ -67,32 +67,32 @@ class TestPostGresql(unittest.TestCase):
     def test_3_load_key(self):
         self.db.savekey('SECRET$6$KEY')
         secretkey = self.db.loadkey()
-        self.assertEqual(secretkey, 'SECRET$6$KEY')
+        self.assertEqual(secretkey, b'SECRET$6$KEY')
 
     def test_4_save_crypto(self):
-        self.db.save_crypto_info("TOP", "SECRET")
+        self.db.save_crypto_info(b"TOP", b"SECRET")
         secretkey = self.db.loadkey()
-        self.assertEqual(secretkey, 'TOP$6$SECRET')
+        self.assertEqual(secretkey, b'TOP$6$SECRET')
         row = self.db.fetch_crypto_info()
-        self.assertEqual(row, ('TOP', 'SECRET'))
+        self.assertEqual(list(map(bytearray, row)),
+                         [bytearray(b'TOP'), bytearray(b'SECRET')])
 
     def test_5_add_node(self):
-        # fuck, saving b"TBONE" has is harder ...
-        innode = ["TBONE", "S3K43T", "example.org", "some note",
-                  ["footag", "bartag"]]
+        innode = [b"TBONE", b"S3K43T", b"example.org", b"some note",
+                  [b"footag", b"bartag"]]
         self.db.add_node(innode)
         outnode = self.db.getnodes([1])[0]
         self.assertEqual(innode[:-1] + [t for t in innode[-1]], outnode[1:])
 
     def test_6_list_nodes(self):
-        ret = self.db.listnodes()
-        self.assertEqual(ret, [1])
-        ret = self.db.listnodes("footag")
-        self.assertEqual(ret, [1])
+        ret1 = self.db.listnodes()
+        self.assertEqual(ret1, [1])
+        ret2 = self.db.listnodes(b"footag")
+        self.assertEqual(ret2, [1])
 
     def test_6a_list_tags(self):
         ret = self.db.listtags()
-        self.assertListEqual(ret, ['footag', 'bartag'])
+        self.assertListEqual(ret, [b'footag', b'bartag'])
 
     def test_6b_get_nodes(self):
         ret = self.db.getnodes([1])
@@ -100,8 +100,8 @@ class TestPostGresql(unittest.TestCase):
         self.assertListEqual(ret, retb)
 
     def test_7_get_or_create_tag(self):
-        s = self.db._get_or_create_tag("SECRET")
-        s1 = self.db._get_or_create_tag("SECRET")
+        s = self.db._get_or_create_tag(b"SECRET")
+        s1 = self.db._get_or_create_tag(b"SECRET")
 
         self.assertEqual(s, s1)