Added ability to add host to blocked table if visit count is exceeded

lib/db.js

@@ -19,6 +19,10 @@ const logger = winston.createLogger({
  *          on exit.
  */
 
+/**
+ * Log a message at level ERROR and exit
+ * @param {string} message
+ */
 function logErrorAndExit(message) {
     logger.error(message)
     process.exit(1)
@@ -92,6 +96,19 @@ function addHostToTable(hostname) {
     })
 }
 
+/**
+ * Add a host to the blocked table, hit after visit count is exceeded.
+ * @param {string} hostname Host to add to blocked table
+ */
+function addHostToBlockedTable(hostname) {
+    db.serialize(() => {
+        db.run("insert into hosts (hostname) values (?)", [hostname], (err, row) => {
+            if (err !== null) {
+                logger.info("Attempted to visit blocked page " + hostname)
+            }
+        })
+    });
+}
 /**
  * Either initialize a row in the database, or update the visit count for a row.
  * @param {string} hostname Host that was visited
@@ -116,6 +133,7 @@ exports.visitHost = function (hostname, callbackFn) {
             } else {
                 // Do the blocking, make this configurable.
                 if (row.visitcount >= 2) {
+                    addHostToBlockedTable(hostname)
                     callbackFn(false);
                 } else {
                     updateVisitCount(hostname);

test/test_proxy.py

@@ -13,7 +13,15 @@ def test_proxy_block_page():
         requests.get("http://httpbin.org", proxies={"http": http_proxy})
     assert requests.get("http://httpbin.org", proxies={"http": http_proxy}).status_code == 403
 
+    # TODO - Use a sqlite driver here to check that httpbin.org is in the blocked table.
+    #        Or mock the sqlite client.
+
 def test_blocked_file():
+    """
+    Tests that the host configured to be blocked is successfully blocked.
+
+    TODO - Mock the blocked file to be used by the proxy to not have to hardcode the host.
+    """
     resp = requests.get("http://wikipedia.org", proxies={"http": http_proxy})
     assert resp.status_code == 403