pgsqlpot 2.0.0__tar.gz → 2.0.2__tar.gz

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/CHANGELOG.md

@@ -5,6 +5,36 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.0.2]
+
+### Added in version 2.0.2
+
+* Nothing
+
+### Changed in version 2.0.2
+
+* Increased the version number
+* The `restart` command wasn't working correctly on Windows due to a race
+  condition. Fixed.
+* Fixed a problem in the MySQL plugin that made it unresponsive under high
+  traffic
+
+## [2.0.1]
+
+### Added in version 2.0.1
+
+* Nothing
+
+### Changed in version 2.0.1
+
+* Increased the version number
+* The `datadog`, `discord`, `nlcvapi`, and `telegram` plugins now use a secure
+  connection (HTTPS) by default
+* The `elastic` plugin now warns if the `ssl` is set while certificate
+  verification (`verify_certs`) is off
+* The `couch` plugin now uses authentication mechanism that does not pass the
+  username and password in the URL
+
 ## [2.0.0]
 
 ### Added in version 2.0.0

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/MANIFEST.in

@@ -1,9 +1,8 @@
-include honeypot.py
-include setup.cfg
+include CHANGELOG.md
 include LICENSE
 include README.md
-include CHANGELOG.md
-graft pgsqlpot
+include honeypot.py
 graft core
 graft output_plugins
+graft pgsqlpot
 prune pgsqlpot/data/test/develop

{pgsqlpot-2.0.0/pgsqlpot.egg-info → pgsqlpot-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pgsqlpot
-Version: 2.0.0
+Version: 2.0.2
 Summary: A PostgeSQL Honeypot
 Home-page: https://gitlab.com/bontchev/pgsqlpot
 Author: Vesselin Bontchev
@@ -43,11 +43,13 @@ Requires-Dist: twisted>=21; python_version >= "3"
 Provides-Extra: couchdb
 Requires-Dist: couchdb; extra == "couchdb"
 Provides-Extra: datadog
+Requires-Dist: certifi; extra == "datadog"
 Requires-Dist: cryptography<=2.8; python_version < "3" and extra == "datadog"
 Requires-Dist: pyOpenSSL<=18.0.0; python_version < "3" and extra == "datadog"
 Requires-Dist: cryptography; python_version >= "3" and extra == "datadog"
 Requires-Dist: pyOpenSSL; python_version >= "3" and extra == "datadog"
 Provides-Extra: discord
+Requires-Dist: certifi; extra == "discord"
 Provides-Extra: elastic
 Requires-Dist: elasticsearch<=7.13; python_version < "3" and extra == "elastic"
 Requires-Dist: numpy<=1.16.6; python_version < "3" and extra == "elastic"
@@ -71,6 +73,7 @@ Provides-Extra: mysql
 Requires-Dist: PyMySQL; python_version < "3" and extra == "mysql"
 Requires-Dist: mysqlclient>=1.3.12; python_version >= "3" and extra == "mysql"
 Provides-Extra: nlcvapi
+Requires-Dist: certifi; extra == "nlcvapi"
 Requires-Dist: pyOpenSSL<=18.0.0; python_version < "3" and extra == "nlcvapi"
 Requires-Dist: pyOpenSSL; python_version >= "3" and extra == "nlcvapi"
 Provides-Extra: postgres
@@ -87,12 +90,14 @@ Requires-Dist: slack-sdk; python_version >= "3" and extra == "slack"
 Provides-Extra: socketlog
 Provides-Extra: sqlite
 Provides-Extra: telegram
+Requires-Dist: certifi; extra == "telegram"
 Provides-Extra: textlog
 Provides-Extra: xmpp
 Requires-Dist: xmpppy>=0.7.3; extra == "xmpp"
 Provides-Extra: all
 Requires-Dist: Automat<20; python_version < "3" and extra == "all"
 Requires-Dist: PyMySQL; python_version < "3" and extra == "all"
+Requires-Dist: certifi; extra == "all"
 Requires-Dist: confluent-kafka; python_version >= "3" and extra == "all"
 Requires-Dist: confluent-kafka<1.0; python_version < "3" and extra == "all"
 Requires-Dist: couchdb; extra == "all"

pgsqlpot-2.0.2/core/httpclient.py

@@ -0,0 +1,71 @@
+from __future__ import absolute_import
+
+from io import open
+
+from twisted.internet.ssl import Certificate, trustRootFromCertificates
+from twisted.python.log import msg
+from twisted.web.client import Agent, BrowserLikePolicyForHTTPS
+
+
+class _LegacyWebClientContextFactory(object):
+    @staticmethod
+    def build():
+        from twisted.internet.ssl import ClientContextFactory
+
+        class WebClientContextFactory(ClientContextFactory):
+            def getContext(self, hostname, port):
+                return ClientContextFactory.getContext(self)
+
+        return WebClientContextFactory()
+
+
+def _load_trust_root_from_pem_bundle(pem_path):
+    with open(pem_path, 'r', encoding='utf-8') as fh:
+        content = fh.read()
+    certs = []
+    end_marker = '-----END CERTIFICATE-----'
+    for chunk in content.split(end_marker):
+        chunk = chunk.strip()
+        if not chunk:
+            continue
+        pem = chunk + '\n' + end_marker + '\n'
+        certs.append(Certificate.loadPEM(pem))
+    if not certs:
+        return None
+    return trustRootFromCertificates(certs)
+
+
+def _trust_root(ca_certs):
+    if ca_certs:
+        try:
+            return _load_trust_root_from_pem_bundle(ca_certs)
+        except Exception as e:
+            msg('Failed to load CA bundle {}: {}'.format(ca_certs, e))
+            return None
+
+    try:
+        import certifi
+        return _load_trust_root_from_pem_bundle(certifi.where())
+    except Exception:
+        return None
+
+
+def create_http_agent(reactor, pool, plugin_name, verify_tls=True, ca_certs=None):
+    """
+    Build a Twisted Agent with sane HTTPS defaults.
+    verify_tls=True uses BrowserLikePolicyForHTTPS for cert + hostname checks.
+    verify_tls=False falls back to a legacy context factory for compatibility.
+    """
+    if verify_tls:
+        return Agent(
+            reactor,
+            contextFactory=BrowserLikePolicyForHTTPS(trustRoot=_trust_root(ca_certs)),
+            pool=pool
+        )
+
+    msg('{}: TLS certificate verification is disabled.'.format(plugin_name))
+    return Agent(
+        reactor,
+        contextFactory=_LegacyWebClientContextFactory.build(),
+        pool=pool
+    )

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/core/tools.py

@@ -15,12 +15,14 @@ from twisted.python.log import msg
 try:
     from urllib.request import urlopen
 except ImportError:
-    from urllib import urlopen
+    from urllib2 import urlopen # type: ignore
 
 
 if version_info[0] >= 3:
     def decode(x):
-        return x.decode('utf-8', errors='ignore')
+        if isinstance(x, bytes):
+            return x.decode('utf-8', errors='ignore')
+        return x
     def encode(x):
         return x.encode()
     def ord(x):
@@ -85,10 +87,7 @@ def stop_plugins(cfg):
 
 def get_public_ip(ip_reporter):
     try:
-        if version_info[0] < 3:
-            return urlopen(ip_reporter).read().decode('latin1', errors='replace').encode('utf-8')
-        else:
-            return decode(urlopen(ip_reporter).read())
+        return decode(urlopen(ip_reporter, timeout=10).read())
     except:
         return None
 

{pgsqlpot-2.0.0/pgsqlpot → pgsqlpot-2.0.2}/honeypot.py

@@ -23,7 +23,7 @@ from twisted.internet.reactor import listenTCP, run
 from twisted.python.log import msg
 
 
-__VERSION__ = '2.0.0'
+__VERSION__ = '2.0.2'
 __description__ = 'A PostgreSQL Honeypot'
 __license__ = 'GPLv3'
 __uri__ = 'https://gitlab.com/bontchev/pgsqlpot'

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/couch.py

@@ -13,6 +13,7 @@ from twisted.python.log import msg
 class Output(output.Output):
 
     def start(self):
+        scheme = CONFIG.get('output_couch', 'scheme', fallback='http')
         host = CONFIG.get('output_couch', 'host', fallback='localhost')
         port = CONFIG.getint('output_couch', 'port', fallback=5984)
         username = CONFIG.get('output_couch', 'username', fallback='pgsqlpot', raw=True)
@@ -20,7 +21,10 @@ class Output(output.Output):
         db_name = CONFIG.get('output_couch', 'database', fallback='pgsqlpot')
 
         try:
-            couchserver = Server('http://{}:{}@{}:{}'.format(username, password, host, port))
+            base_url = '{}://{}:{}'.format(scheme, host, port)
+            couchserver = Server(base_url)
+            if username:
+                couchserver.resource.credentials = (username, password)
 
             if db_name in couchserver:
                 self.couch_db = couchserver[db_name]

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/datadog.py

@@ -10,18 +10,13 @@ from platform import node
 
 from core import output
 from core.config import CONFIG
+from core.httpclient import create_http_agent
 from core.tools import to_bytes
 
 from twisted.internet import reactor
 from twisted.python.log import msg
 from twisted.web import client, http_headers
 from twisted.web.client import FileBodyProducer
-from twisted.internet.ssl import ClientContextFactory
-
-
-class WebClientContextFactory(ClientContextFactory):
-    def getContext(self, hostname, port):
-        return ClientContextFactory.getContext(self)
 
 
 class QuietHTTP11ClientFactory(client._HTTP11ClientFactory):
@@ -33,16 +28,17 @@ class Output(output.Output):
         self.url = CONFIG.get('output_datadog', 'url')
         self.api_key = CONFIG.get('output_datadog', 'api_key', fallback='')
         if len(self.api_key) == 0:
-            msg('Datadog output module: API key is not defined.')
+            msg('output_datadog: API key is not defined.')
         self.ddsource = CONFIG.get('output_datadog', 'ddsource', fallback='pgsqlpot')
         self.ddtags = CONFIG.get('output_datadog', 'ddtags', fallback='env:dev')
         self.service = CONFIG.get('output_datadog', 'service', fallback='honeypot')
         self.hostname = CONFIG.get('output_datadog', 'hostname', fallback=node())
+        verify_tls = CONFIG.getboolean('output_datadog', 'verify_tls', fallback=True)
+        ca_certs = CONFIG.get('output_datadog', 'ca_certs', fallback=None)
 
-        contextFactory = WebClientContextFactory()
         myQuietPool = client.HTTPConnectionPool(reactor)
         myQuietPool._factory = QuietHTTP11ClientFactory
-        self.agent = client.Agent(reactor, contextFactory=contextFactory, pool=myQuietPool)
+        self.agent = create_http_agent(reactor, myQuietPool, 'output_datadog', verify_tls, ca_certs)
 
     def stop(self):
         pass

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/discord.py

@@ -13,14 +13,13 @@ from time import gmtime, strftime, time
 
 from core import output
 from core.config import CONFIG
+from core.httpclient import create_http_agent
 from core.tools import decode, to_bytes
 
 from twisted.internet import reactor
-from twisted.internet.ssl import ClientContextFactory
 from twisted.internet.task import deferLater
 from twisted.python.log import msg
 from twisted.web.client import (
-    Agent,
     FileBodyProducer,
     HTTPConnectionPool,
     _HTTP11ClientFactory,
@@ -29,11 +28,6 @@ from twisted.web.client import (
 from twisted.web.http_headers import Headers
 
 
-class WebClientContextFactory(ClientContextFactory):
-    def getContext(self, hostname, port):
-        return ClientContextFactory.getContext(self)
-
-
 class QuietHTTP11ClientFactory(_HTTP11ClientFactory):
     noisy = False
 
@@ -43,10 +37,11 @@ class Output(output.Output):
     def start(self):
         self.url = to_bytes(CONFIG.get('output_discord', 'url'))
         self.delay = CONFIG.getfloat('output_discord', 'delay', fallback=2.0)
-        contextFactory = WebClientContextFactory()
+        verify_tls = CONFIG.getboolean('output_discord', 'verify_tls', fallback=True)
+        ca_certs = CONFIG.get('output_discord', 'ca_certs', fallback=None)
         pool = HTTPConnectionPool(reactor)
         pool._factory = QuietHTTP11ClientFactory
-        self.agent = Agent(reactor, contextFactory=contextFactory, pool=pool)
+        self.agent = create_http_agent(reactor, pool, 'output_discord', verify_tls, ca_certs)
         self.last_sent = 0
         self.requests_list = []
 

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/elastic.py

@@ -44,6 +44,8 @@ class Output(output.Output):
             kwargs['verify_certs'] = verify_certs
             if verify_certs and ca_certs:
                 kwargs['ca_certs'] = ca_certs
+            elif not verify_certs:
+                msg('output_elastic: SSL is enabled but certificate verification is disabled.')
 
         # Create client
         self.es = Elasticsearch(hosts=hosts, **kwargs)

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/mysql.py

@@ -24,9 +24,16 @@ except ImportError:
         from _mysql_exceptions import (Error, OperationalError) # type: ignore
 
 from twisted.enterprise.adbapi import ConnectionPool
-from twisted.python.compat import reraise
 from twisted.python.log import msg
 
+if version_info[0] >= 3:
+    def _reraise(tp, value, tb):
+        raise value.with_traceback(tb)
+else:
+    exec("""def _reraise(tp, value, tb):
+    raise tp, value, tb
+""")
+
 
 class ReconnectingConnectionPool(ConnectionPool):
     """
@@ -42,14 +49,27 @@ class ReconnectingConnectionPool(ConnectionPool):
     def _runInteraction(self, interaction, *args, **kw):
 
         def rerise_exception(conn):
-            _, excValue, excTraceback = exc_info()
+            tp, value, tb = exc_info()
             try:
                 conn.rollback()
             except Exception:
                 msg('Rollback failed')
-            reraise(excValue, excTraceback)
+            _reraise(tp, value, tb)
+
+        conn = self.connect()
+
+        # Ping the connection before use.  This transparently handles stale
+        # connections closed by MySQL's wait_timeout (default 8 hours): if the
+        # server closed our idle connection, ping(True) reconnects immediately.
+        # If MySQL is genuinely unreachable, ping(True) raises OperationalError
+        # within connect_timeout seconds, which is caught below and re-raised
+        # so the Deferred errback fires and the worker thread is freed promptly.
+        try:
+            conn.ping(True)
+        except Exception:
+            self.disconnect(conn)
+            raise
 
-        conn = self.connectionFactory(self)
         trans = self.transactionFactory(self, conn)
         try:
             result = interaction(trans, *args, **kw)
@@ -57,12 +77,10 @@ class ReconnectingConnectionPool(ConnectionPool):
             conn.commit()
             return result
         except OperationalError as e:
-            if e.args[0] not in (2003, 2006, 2013):
-                rerise_exception(conn)
-            else:
+            if e.args[0] in (2003, 2006, 2013):
                 conn = self.connections.get(self.threadID())
                 self.disconnect(conn)
-                return ConnectionPool._runInteraction(self, interaction, *args, **kw)
+            rerise_exception(conn)
         except Exception:
             rerise_exception(conn)
 
@@ -73,6 +91,35 @@ class Output(output.Output):
         if self.debug:
             msg(message)
 
+    def _enqueue(self, label, d):
+        """
+        Wrap a runInteraction() deferred with two targeted diagnostics:
+
+        1. Pending counter + threshold warning: if the worker threads are stuck
+           the counter climbs without ever coming back down.  A warning fires
+           once per threshold crossing so the log stays quiet under normal load
+           but shouts when something is wrong.
+
+        2. Always-on errback: any DB error that would previously have been
+           silently swallowed now produces an explicit log line.
+        """
+        self._pending += 1
+
+        def on_success(result):
+            self._pending -= 1
+            return result
+
+        def on_failure(failure):
+            self._pending -= 1
+            msg('output_mysql: runInteraction {} failed: {}'.format(label, failure))
+            return None
+
+        d.addCallback(on_success)
+        d.addErrback(on_failure)
+
+        if self._pending > self._pending_warn_threshold:
+            msg('output_mysql: WARNING - {} interactions pending, worker thread may be stuck'.format(self._pending))
+
     def start(self):
         host = CONFIG.get('output_mysql', 'host', fallback='localhost')
         database = CONFIG.get('output_mysql', 'database', fallback='pgsqlpot')
@@ -82,6 +129,13 @@ class Output(output.Output):
 
         self.debug = CONFIG.getboolean('output_mysql', 'debug', fallback=False)
         self.geoip = CONFIG.getboolean('output_mysql', 'geoip', fallback=True)
+        self._pending = 0
+        self._pending_warn_threshold = CONFIG.getint('output_mysql', 'pending_warn_threshold', fallback=100)
+
+        connect_timeout = CONFIG.getint('output_mysql', 'connect_timeout', fallback=10)
+        read_timeout    = CONFIG.getint('output_mysql', 'read_timeout',    fallback=30)
+        write_timeout   = CONFIG.getint('output_mysql', 'write_timeout',   fallback=30)
+        cp_max = CONFIG.getint('output_mysql', 'cp_max', fallback=5)
 
         try:
             self.dbh = ReconnectingConnectionPool(
@@ -93,11 +147,14 @@ class Output(output.Output):
                 port=port,
                 charset='utf8',
                 use_unicode=True,
+                connect_timeout=connect_timeout,
+                read_timeout=read_timeout,
+                write_timeout=write_timeout,
                 cp_min=1,
-                cp_max=1
+                cp_max=cp_max,
             )
         except Error as e:
-            self.local_log('output_mysql: MySQL Error {}: "{}"'.format(e.args[0], e.args[1]))
+            msg('output_mysql: MySQL Error {}: "{}"'.format(e.args[0], e.args[1]))
 
         if self.geoip:
             geoipdb_city_path = CONFIG.get('output_mysql', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')
@@ -106,13 +163,14 @@ class Output(output.Output):
                 self.reader_city = Reader(geoipdb_city_path)
             except Exception:
                 self.reader_city = None
-                self.local_log('Failed to open City GeoIP database {}'.format(geoipdb_city_path))
+                msg('output_mysql: Failed to open City GeoIP database {}'.format(geoipdb_city_path))
 
             try:
                 self.reader_asn = Reader(geoipdb_asn_path)
             except Exception:
                 self.reader_asn = None
-                self.local_log('Failed to open ASN GeoIP database {}'.format(geoipdb_asn_path))
+                msg('output_mysql: Failed to open ASN GeoIP database {}'.format(geoipdb_asn_path))
+
 
     def stop(self):
         if self.geoip:
@@ -120,7 +178,6 @@ class Output(output.Output):
                 self.reader_city.close()
             if self.reader_asn is not None:
                 self.reader_asn.close()
-
     def write(self, event):
         """
         TODO: Check if the type (date, datetime or timestamp) of columns is appropriate for your needs and timezone
@@ -130,7 +187,8 @@ class Output(output.Output):
         and back from UTC to the current time zone for retrieval.
         (This does not occur for other types such as DATETIME.)"
         """
-        self.dbh.runInteraction(self.connect_event, event)
+        self._enqueue('connect_event', self.dbh.runInteraction(self.connect_event, event))
+
 
     def simple_query(self, txn, sql, args):
         if self.debug:
@@ -145,22 +203,19 @@ class Output(output.Output):
                 txn.execute(sql)
             result = txn.fetchall()
         except Exception as e:
-            self.local_log('output_mysql: MySQL Error: {}'.format(e))
+            msg('output_mysql: MySQL Error: {}'.format(e))
             result = None
         return result
 
     def get_id(self, txn, table, column, entry):
+        # INSERT IGNORE silently skips the insert when a UNIQUE constraint would
+        # be violated, so the subsequent SELECT always finds exactly one row
+        # regardless of whether a concurrent call already inserted it.
+        self.simple_query(txn, "INSERT IGNORE INTO `{}` (`{}`) VALUES (%s)".format(table, column), (entry, ))
         r = self.simple_query(txn, "SELECT `id` FROM `{}` WHERE `{}` = %s".format(table, column), (entry, ))
         if r:
-            id = r[0][0]
-        else:
-            self.simple_query(txn, "INSERT INTO `{}` (`{}`) VALUES (%s)".format(table, column), (entry, ))
-            r = self.simple_query(txn, 'SELECT LAST_INSERT_ID()', ())
-            if r:
-                id = int(r[0][0])
-            else:
-                id = 0
-        return id
+            return r[0][0]
+        return 0
 
     def connect_event(self, txn, event):
         remote_ip = event['src_ip']
@@ -197,14 +252,15 @@ class Output(output.Output):
 
         if self.geoip:
             country, country_code, city, org, asn_num = geolocate(remote_ip, self.reader_city, self.reader_asn)
+            # INSERT IGNORE rather than ON DUPLICATE KEY UPDATE: geolocation
+            # data for an IP rarely changes, so skipping the update on
+            # subsequent hits is acceptable.  More importantly, ON DUPLICATE
+            # KEY UPDATE takes an exclusive row lock on the existing row,
+            # causing InnoDB lock contention when multiple threads process
+            # connections from the same IP simultaneously.
+            # INSERT IGNORE avoids that lock entirely.
             self.simple_query(txn, """
-                INSERT INTO `geolocation` (`ip`, `country_name`, `country_iso_code`, `city_name`, `org`, `org_asn`)
+                INSERT IGNORE INTO `geolocation` (`ip`, `country_name`, `country_iso_code`, `city_name`, `org`, `org_asn`)
                 VALUES (%s, %s, %s, %s, %s, %s)
-                ON DUPLICATE KEY UPDATE
-                    `country_name` = %s,
-                    `country_iso_code` = %s,
-                    `city_name` = %s,
-                    `org` = %s,
-                    `org_asn` = %s
                 """,
-                (remote_ip, country, country_code, city, org, asn_num, country, country_code, city, org, asn_num, ))
+                (remote_ip, country, country_code, city, org, asn_num, ))

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/nlcvapi.py

@@ -8,15 +8,14 @@ from json import dumps, loads
 
 from core import output
 from core.config import CONFIG
+from core.httpclient import create_http_agent
 from core.tools import decode, geolocate, to_bytes
 
 from geoip2.database import Reader
 
 from twisted.internet import reactor
-from twisted.internet.ssl import ClientContextFactory
 from twisted.python.log import msg
 from twisted.web.client import (
-    Agent,
     FileBodyProducer,
     HTTPConnectionPool,
     _HTTP11ClientFactory,
@@ -25,12 +24,6 @@ from twisted.web.client import (
 from twisted.web.http_headers import Headers
 
 
-class WebClientContextFactory(ClientContextFactory):
-
-    def getContext(self, hostname, port):
-        return ClientContextFactory.getContext(self)
-
-
 class QuietHTTP11ClientFactory(_HTTP11ClientFactory):
     noisy = False
 
@@ -40,10 +33,11 @@ class Output(output.Output):
     def start(self):
         self.host = CONFIG.get('output_nlcvapi', 'host', fallback='https://api.nlcv.bas.bg/v1.0/honeypot')
         self.geoip = CONFIG.getboolean('output_nlcvapi', 'geoip', fallback=True)
-        contextFactory = WebClientContextFactory()
+        verify_tls = CONFIG.getboolean('output_nlcvapi', 'verify_tls', fallback=True)
+        ca_certs = CONFIG.get('output_nlcvapi', 'ca_certs', fallback=None)
         pool = HTTPConnectionPool(reactor)
         pool._factory = QuietHTTP11ClientFactory
-        self.agent = Agent(reactor, contextFactory=contextFactory, pool=pool)
+        self.agent = create_http_agent(reactor, pool, 'output_nlcvapi', verify_tls, ca_certs)
 
         if self.geoip:
             geoipdb_city_path = CONFIG.get('output_nlcvapi', 'geoip_citydb', fallback='data/GeoLite2-City.mmdb')

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/postgres.py

@@ -74,24 +74,18 @@ class Output(output.Output):
         return result
 
     def get_id(self, txn, table, column, entry):
+        # ON CONFLICT ... DO UPDATE is a deliberate no-op that makes RETURNING
+        # yield the existing row's id even when the INSERT is skipped due to a
+        # unique-constraint conflict.  DO NOTHING would leave RETURNING empty.
         r = self.simple_query(
             txn,
-            "SELECT id FROM {} WHERE {} = %s".format(table, column),
-            (entry, )
-        )
+            "INSERT INTO {} ({}) VALUES (%s) "
+            "ON CONFLICT ({}) DO UPDATE SET {} = EXCLUDED.{} "
+            "RETURNING id".format(table, column, column, column, column),
+            (entry, ))
         if r:
-            id = int(r[0])
-        else:
-            r = self.simple_query(
-                txn,
-                "INSERT INTO {} ({}) VALUES (%s) RETURNING id".format(table, column),
-                (entry, )
-            )
-            if r:
-                id = int(r[0])
-            else:
-                id = 0
-        return id
+            return int(r[0])
+        return 0
 
     def connect_event(self, txn, event):
         remote_ip = event['src_ip']

{pgsqlpot-2.0.0 → pgsqlpot-2.0.2}/output_plugins/sqlite.py

@@ -15,7 +15,7 @@ from twisted.python.log import msg
 class Output(output.Output):
 
     def start(self):
-        db_name = CONFIG.get('output_sqlite', 'db_file', fallback='data/pgsqlpot.db')
+        db_name = CONFIG.get('output_sqlite', 'db_file', fallback='log/pgsqlpot.db')
         self.debug = CONFIG.getboolean('output_sqlite', 'debug', fallback=False)
         self.geoip = CONFIG.getboolean('output_sqlite', 'geoip', fallback=True)
 
@@ -75,17 +75,13 @@ class Output(output.Output):
         return result
 
     def get_id(self, txn, table, column, entry):
+        # INSERT OR IGNORE silently skips the insert when a UNIQUE constraint
+        # would be violated, so the subsequent SELECT always finds one row.
+        self.simple_query(txn, "INSERT OR IGNORE INTO `{}` (`{}`) VALUES (?)".format(table, column), (entry, ))
         r = self.simple_query(txn, "SELECT `id` FROM `{}` WHERE `{}` = ?".format(table, column), (entry, ))
         if r:
-            id = r[0][0]
-        else:
-            self.simple_query(txn, "INSERT INTO `{}` (`{}`) VALUES (?)".format(table, column), (entry, ))
-            r = self.simple_query(txn, 'SELECT LAST_INSERT_ROWID()', ())
-            if r:
-                id = int(r[0][0])
-            else:
-                id = 0
-        return id
+            return r[0][0]
+        return 0
 
     def connect_event(self, txn, event):
         remote_ip = event['src_ip']
@@ -100,7 +96,7 @@ class Output(output.Output):
             """,
             (event['session'], event['unixtime'], operation_id, remote_ip, event['src_port'],
              event['dst_ip'], event['dst_port'], sensor_id, ))
-        
+
         if event['operation'].lower() == 'login':
             usr_id = self.get_id(txn, 'usernames', 'username', event['username'])
             pwd_id = self.get_id(txn, 'passwords', 'password', event['password'])