Refactored whois-related logic

Author: sarukuku

db.js

@@ -4,33 +4,48 @@ const argv = require('yargs').argv
 const sqlite3 = require('sqlite3').verbose()
 let db
 
-exports.openOrCreate = function () {
-  db = new sqlite3.Database('whois-cache.sqlite3', initDatabase)
+exports.openOrCreate = async () => {
+  return new Promise(resolve => {
+    db = new sqlite3.Database('whois-cache.sqlite3', async () => {
+      await initDatabase()
+      resolve()
+    })
+  })
 }
 
-const initDatabase = () => {
-  db.run(`
-    CREATE TABLE IF NOT EXISTS whois (
-      domain TEXT PRIMARY KEY,
-      registrantName TEXT,
-      registrantOrganization TEXT,
-      registrantCountry TEXT
-    )
-  `)
+const initDatabase = async () => {
+  return new Promise(resolve => {
+    return db.run(`
+      CREATE TABLE IF NOT EXISTS whois (
+        domain TEXT PRIMARY KEY,
+        registrantName TEXT,
+        registrantOrganization TEXT,
+        registrantCountry TEXT
+      )
+    `, () => {
+      resolve()
+    })
+  })
 }
 
-exports.insertDomain = (whoisObj) => {
-  db.run(`
-    REPLACE INTO whois VALUES (
-      '${whoisObj.domain}',
-      '${whoisObj.registrantName}',
-      '${whoisObj.registrantOrganization}',
-      '${whoisObj.registrantCountry}'
+exports.insertDomain = async whoisObj => {
+  return new Promise(resolve => {
+    let statement = db.prepare('REPLACE INTO whois VALUES (?,?,?,?)')
+
+    statement.run(
+      whoisObj.domain,
+      whoisObj.registrantName,
+      whoisObj.registrantOrganization,
+      whoisObj.registrantCountry
     )
-  `)
+
+    statement.finalize(() => {
+      resolve()
+    })
+  })
 }
 
-exports.readDomain = (domain) => {
+exports.readDomain = domain => {
   return new Promise(resolve => {
     db.get(`SELECT * FROM whois WHERE domain = '${domain}'`, (err, row) => {
       if (err) {

index.js

@@ -3,7 +3,6 @@
 const utils = require('./utils.js')
 const whoisDb = require('./db.js')
 const argv = require('yargs').argv
-const async = require('async')
 const puppeteer = require('puppeteer')
 const devices = require('puppeteer/DeviceDescriptors')
 const iPadPro = devices['iPad Pro landscape'];
@@ -14,10 +13,6 @@ const iPadPro = devices['iPad Pro landscape'];
     return
   }
 
-  // Open on init whois sqlite3 database.
-  utils.log('Opening / initializing whois cache database...')
-  whoisDb.openOrCreate()
-
   utils.startTime('Test duration')
   utils.log('Starting the test...')
   const browser = await puppeteer.launch()
@@ -87,40 +82,16 @@ const iPadPro = devices['iPad Pro landscape'];
 
   // If -l (=long) is set --> Collect whois data
   if (argv.l) {
+    // Open on init whois sqlite3 database.
+    utils.log('Opening / initializing whois cache database...')
+    await whoisDb.openOrCreate()
+
     utils.log('Collecting whois data for cross-domains...')
-    let collectedWhoisData = []
-    for (let key in data.resources) {
-      if (data.resources.hasOwnProperty(key)) {
-        utils.log(`Collecting whois data of ${key} sources...`)
-        await new Promise(resolve => {
-          // Limit concurrent whois data.resources to 30
-          async.mapLimit(data.resources[key].requests, 30, async (urlObj) => {
-            let rootDomain = utils.getRootDomain(urlObj.url)
-
-            // Fetch whois data only for cross-origins
-            if (urlObj.crossOrigin && !utils.whoisDataValid(collectedWhoisData[rootDomain])) {
-              let whoisData = await utils.getOwnerData(rootDomain)
-              if (utils.whoisDataValid(whoisData)) {
-                collectedWhoisData[rootDomain] = whoisData
-              }
-            }
-
-            if (urlObj.crossOrigin) {
-              urlObj.ownerData = collectedWhoisData[rootDomain]
-            }
-          }, (err, results) => {
-            if (err) {
-              utils.log(err)
-            }
-            resolve()
-          })
-        })
-      }
-    }
-  }
+    data = await utils.collectWhoisData(data)
 
-  // Close the database as we don't need it anymore.
-  whoisDb.close()
+    // Close the database as we don't need it anymore.
+    whoisDb.close()
+  }
 
   // Print results to stdout based on selected output.
   utils.log('Printing results...')

utils.js

@@ -4,6 +4,7 @@ const { spawn } = require('child_process')
 const argv = require('yargs').argv
 const URL = require('url-parse')
 const parseDomain = require('parse-domain')
+const async = require('async')
 const whoisDb = require('./db.js')
 
 function getWhoisData (domain) {
@@ -160,7 +161,7 @@ exports.getOwnerData = async function (domain) {
 
   // Save whois data to DB only if it's valid
   if (exports.whoisDataValid(data)) {
-    whoisDb.insertDomain(data)
+    await whoisDb.insertDomain(data)
   }
 
   return data
@@ -228,9 +229,9 @@ Cross-origin percentage: ${data.resources[key].crossOriginPercentage}%`)
             console.log(`
   URL: ${exports.truncate.apply(requestObj.url, [100, false])}
   Owner data:
-    Registrant Name: ${(requestObj.ownerData && requestObj.ownerData.registrantName) ? requestObj.ownerData.registrantName : ''}
-    Registrant Organization: ${(requestObj.ownerData && requestObj.ownerData.registrantOrganization) ? requestObj.ownerData.registrantOrganization : ''}
-    Registrant Country: ${(requestObj.ownerData && requestObj.ownerData.registrantCountry) ? requestObj.ownerData.registrantCountry : ''}`)
+    Registrant Name: ${(requestObj.whoisData && requestObj.whoisData.registrantName) ? requestObj.whoisData.registrantName : ''}
+    Registrant Organization: ${(requestObj.whoisData && requestObj.whoisData.registrantOrganization) ? requestObj.whoisData.registrantOrganization : ''}
+    Registrant Country: ${(requestObj.whoisData && requestObj.whoisData.registrantCountry) ? requestObj.whoisData.registrantCountry : ''}`)
           }
         })
       }
@@ -312,3 +313,35 @@ exports.startTime = id => {
 exports.endTime = id => {
   !argv.silent && console.timeEnd(id)
 }
+
+exports.collectWhoisData = async data => {
+  let collectedWhoisData = []
+  for (let key in data.resources) {
+    exports.log(`Collecting whois data of ${key} sources...`)
+    await new Promise(resolve => {
+      // Limit concurrent whois queries to 30
+      async.mapLimit(data.resources[key].requests, 30, async (requestObj) => {
+        let rootDomain = exports.getRootDomain(requestObj.url)
+
+        // Fetch whois data only for cross-origins
+        if (requestObj.crossOrigin && !exports.whoisDataValid(collectedWhoisData[rootDomain])) {
+          let whoisData = await exports.getOwnerData(rootDomain)
+          if (exports.whoisDataValid(whoisData)) {
+            collectedWhoisData[rootDomain] = whoisData
+          }
+        }
+
+        if (requestObj.crossOrigin) {
+          requestObj.whoisData = collectedWhoisData[rootDomain]
+        }
+      }, (err, results) => {
+        if (err) {
+          exports.log(err)
+        }
+        resolve()
+      })
+    })
+  }
+
+  return data
+}