From 5a40a15f9c1e7d891448d9a8c7af433d90a12086 Mon Sep 17 00:00:00 2001
From: jaketothepast <windlejacob12@gmail.com>
Date: Wed, 13 Feb 2019 16:34:56 -0500
Subject: [PATCH] Added ability to add host to blocked table if visit count is
 exceeded

---
 lib/db.js          | 18 ++++++++++++++++++
 test/test_proxy.py |  8 ++++++++
 2 files changed, 26 insertions(+)

diff --git a/lib/db.js b/lib/db.js
index 9dcd3d3..4a6153c 100644
--- a/lib/db.js
+++ b/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);
diff --git a/test/test_proxy.py b/test/test_proxy.py
index a781e39..25ec0dd 100644
--- a/test/test_proxy.py
+++ b/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
     
\ No newline at end of file