From f8aa81be8de1515e5f85f20afed5a79ab361a226 Mon Sep 17 00:00:00 2001 From: jaketothepast Date: Sun, 27 Jan 2019 21:29:10 -0500 Subject: [PATCH] Implement blocking functionality after certain count of visits, added tests --- index.js | 43 ++++++++++++++++++++++++------- lib/db.js | 17 ++++++++---- package-lock.json | 64 ++++++++++++++++++---------------------------- test/test_proxy.py | 5 ++++ 4 files changed, 76 insertions(+), 53 deletions(-) diff --git a/index.js b/index.js index da48cfe..ea30b6b 100644 --- a/index.js +++ b/index.js @@ -15,7 +15,7 @@ const logger = winston.createLogger({ /** * Connection handler for proxy, forward request and then pipe response * back to the client. - * + * * @param {IncomingMessage} request The request from the client * @param {OutgoingMessage} response The response to write back to the client */ @@ -28,23 +28,48 @@ const processRequest = (request, response) => { // Set up a data handler for the socket connection. request.on("data", (chunk) => { body.push(chunk); - }) - + }); + // Set up an end handler for the socket connection. request.on("end", () => { body = Buffer.concat(body).toString(); logger.info({message: "Requesting: " + req.url}); const hostUrl = new url.URL(req.url); - myDb.visitHost(hostUrl.hostname); - // Complicated line of JS to forward request, and pipe it to the - // response object. - // - // Uses closures to ensure that both req and res are what I want. - req.pipe(http.request(req.url, (resp) => {resp.pipe(res)})); + + // Do the work of actually updating our db with the visit. + // Val in the callback will be false if blocked. + myDb.visitHost(hostUrl.hostname, (val) => { + if (val === false) { + // blocked. + console.log("Blocked you fool.") + sendBlockPage(res); + } else { + // Not blocked + req.pipe(http.request(req.url, (resp) => { + resp.pipe(res); + })) + } + }); }); } +/** + * Send a block page back to the client. + * @param {OutgoingMessage} res Response to the client + */ +function sendBlockPage(res) { + res.writeHead(403, { + 'Content-Type': 'text/html', + 'X-Powered-By': 'bacon' + }) + res.write('

BLOCKED BY MY PROXY

'); + res.end(); +} + +/** + * Start the server. + */ function startServer() { // Start the server, set up data and end handlers. var server = http.createServer(processRequest); diff --git a/lib/db.js b/lib/db.js index e6babf8..b67e444 100644 --- a/lib/db.js +++ b/lib/db.js @@ -1,5 +1,5 @@ var sqlite3 = require('sqlite3').verbose(); -var db = new sqlite3.Database("./proxy.db") +var db = new sqlite3.Database("./proxy.db"); /** TODO - MAKE THESE FUNCTIONS THEIR OWN MODULE. VVV * FIXME - need to close the database connection @@ -35,7 +35,7 @@ function updateVisitCount(hostname) { */ function addHostToTable(hostname) { db.serialize(() => { - db.run("insert into hosts (hostname, visitcount) values (?, 0)", [hostname]) + db.run("insert into hosts (hostname, visitcount) values (?, 0)", [hostname]); }) } @@ -43,14 +43,21 @@ function addHostToTable(hostname) { * Either initialize a row in the database, or update the visit count for a row. * @param {string} hostname Host that was visited */ -exports.visitHost = function (hostname) { +exports.visitHost = function (hostname, callbackFn) { return db.serialize(() => { db.get("select * from hosts where hostname = ?", [hostname], (err, row) => { if (row === undefined) { addHostToTable(hostname); + callbackFn(true); } else { - updateVisitCount(hostname); + // Do the blocking, make this configurable. + if (row.visitcount >= 2) { + callbackFn(false); + } else { + updateVisitCount(hostname); + callbackFn(true); + } } - }) + }); }); } diff --git a/package-lock.json b/package-lock.json index 2d449cd..03d3032 100644 --- a/package-lock.json +++ b/package-lock.json @@ -841,11 +841,6 @@ "tweetnacl": "^0.14.3" } }, - "bluebird": { - "version": "3.5.3", - "resolved": "https://registry.npmjs.org/bluebird/-/bluebird-3.5.3.tgz", - "integrity": "sha512-/qKPUQlaW1OyR51WeCPBvRnAlnZFUJkCSG5HzGnuIqhgyJtF+T94lFnn33eiazjRm2LAHVy2guNnaq48X9SJuw==" - }, "body-parser": { "version": "1.18.3", "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.18.3.tgz", @@ -1867,12 +1862,14 @@ "balanced-match": { "version": "1.0.0", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "brace-expansion": { "version": "1.1.11", "bundled": true, "dev": true, + "optional": true, "requires": { "balanced-match": "^1.0.0", "concat-map": "0.0.1" @@ -1887,17 +1884,20 @@ "code-point-at": { "version": "1.1.0", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "concat-map": { "version": "0.0.1", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "console-control-strings": { "version": "1.1.0", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "core-util-is": { "version": "1.0.2", @@ -2014,7 +2014,8 @@ "inherits": { "version": "2.0.3", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "ini": { "version": "1.3.5", @@ -2026,6 +2027,7 @@ "version": "1.0.0", "bundled": true, "dev": true, + "optional": true, "requires": { "number-is-nan": "^1.0.0" } @@ -2040,6 +2042,7 @@ "version": "3.0.4", "bundled": true, "dev": true, + "optional": true, "requires": { "brace-expansion": "^1.1.7" } @@ -2047,12 +2050,14 @@ "minimist": { "version": "0.0.8", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "minipass": { "version": "2.3.5", "bundled": true, "dev": true, + "optional": true, "requires": { "safe-buffer": "^5.1.2", "yallist": "^3.0.0" @@ -2071,6 +2076,7 @@ "version": "0.5.1", "bundled": true, "dev": true, + "optional": true, "requires": { "minimist": "0.0.8" } @@ -2151,7 +2157,8 @@ "number-is-nan": { "version": "1.0.1", "bundled": true, - "dev": true + "dev": true, + "optional": true }, "object-assign": { "version": "4.1.1", @@ -2163,6 +2170,7 @@ "version": "1.4.0", "bundled": true, "dev": true, + "optional": true, "requires": { "wrappy": "1" } @@ -2284,6 +2292,7 @@ "version": "1.0.2", "bundled": true, "dev": true, + "optional": true, "requires": { "code-point-at": "^1.0.0", "is-fullwidth-code-point": "^1.0.0", @@ -4694,21 +4703,11 @@ } } }, - "request-promise": { - "version": "4.2.2", - "resolved": "https://registry.npmjs.org/request-promise/-/request-promise-4.2.2.tgz", - "integrity": "sha1-0epG1lSm7k+O5qT+oQGMIpEZBLQ=", - "requires": { - "bluebird": "^3.5.0", - "request-promise-core": "1.1.1", - "stealthy-require": "^1.1.0", - "tough-cookie": ">=2.3.3" - } - }, "request-promise-core": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/request-promise-core/-/request-promise-core-1.1.1.tgz", "integrity": "sha1-Pu4AssWqgyOc+wTFcA2jb4HNCLY=", + "dev": true, "requires": { "lodash": "^4.13.1" } @@ -5417,21 +5416,6 @@ "integrity": "sha1-BOaSb2YolTVPPdAVIDYzuFcpfiw=", "dev": true }, - "sql-template-strings": { - "version": "2.2.2", - "resolved": "https://registry.npmjs.org/sql-template-strings/-/sql-template-strings-2.2.2.tgz", - "integrity": "sha1-PxFQiiWt384hejBCqdMAwxk7lv8=", - "optional": true - }, - "sqlite": { - "version": "3.0.0", - "resolved": "https://registry.npmjs.org/sqlite/-/sqlite-3.0.0.tgz", - "integrity": "sha512-bGCCf43nnIcVHRXuQfSv0C9khuKvHGlbzzL0dFeXgnsjRS1Vqjs0yUaLPf3Qt+0j0AKUggEmBumGNDNOl4feig==", - "requires": { - "sql-template-strings": "^2.2.2", - "sqlite3": "^4.0.0" - } - }, "sqlite3": { "version": "4.0.6", "resolved": "https://registry.npmjs.org/sqlite3/-/sqlite3-4.0.6.tgz", @@ -5505,7 +5489,8 @@ "stealthy-require": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/stealthy-require/-/stealthy-require-1.1.1.tgz", - "integrity": "sha1-NbCYdbT/SfJqd35QmzCQoyJr8ks=" + "integrity": "sha1-NbCYdbT/SfJqd35QmzCQoyJr8ks=", + "dev": true }, "string-length": { "version": "2.0.0", @@ -5687,6 +5672,7 @@ "version": "2.5.0", "resolved": "https://registry.npmjs.org/tough-cookie/-/tough-cookie-2.5.0.tgz", "integrity": "sha1-zZ+yoKodWhK0c72fuW+j3P9lreI=", + "dev": true, "requires": { "psl": "^1.1.28", "punycode": "^2.1.1" diff --git a/test/test_proxy.py b/test/test_proxy.py index e32ec59..3ae86db 100644 --- a/test/test_proxy.py +++ b/test/test_proxy.py @@ -1,4 +1,5 @@ import requests +import os # TODO - Make this configurable based on config file test section. http_proxy = "http://127.0.0.1:8124" @@ -7,4 +8,8 @@ def test_proxy_basic(): resp = requests.get("http://httpbin.org", proxies={"http": http_proxy}) assert resp.status_code == 200 +def test_proxy_block_page(): + for i in range(2): + requests.get("http://httpbin.org", proxies={"http": http_proxy}) + assert requests.get("http://httpbin.org", proxies={"http": http_proxy}).status_code == 403 \ No newline at end of file