fix: changed threadpool -> malebolgia

src/takajo.nim

@@ -1,6 +1,8 @@
 import algorithm
 import cligen
 import json
+import malebolgia
+import malebolgia / ticketlocks
 import nancy
 import puppy
 import re
@@ -12,7 +14,6 @@ import tables
 import terminal
 import termstyle
 import times
-import threadpool
 import uri
 import os
 import std/enumerate

src/takajopkg/vtDomainLookup.nim

@@ -3,9 +3,7 @@
 # Add categories and SAN info
 # Add output not found to txt file
 
-var vtAPIDomainChannel: Channel[VirusTotalResult] # channel for receiving parallel query results
-
-proc queryDomainAPI(domain:string, headers: httpheaders.HttpHeaders) {.thread.} =
+proc queryDomainAPI(domain:string, headers: httpheaders.HttpHeaders, results: ptr seq[VirusTotalResult]; L: ptr TicketLock) =
     let response = get("https://www.virustotal.com/api/v3/domains/" & encodeUrl(domain), headers)
     var jsonResponse = %* {}
     var singleResultTable = newTable[string, string]()
@@ -35,8 +33,8 @@ proc queryDomainAPI(domain:string, headers: httpheaders.HttpHeaders) {.thread.}
         singleResultTable["SSL-ValidUntil"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "not_after"])
         singleResultTable["SSL-Issuer"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "issuer", "O"])
         singleResultTable["SSL-IssuerCountry"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "issuer", "C"])
-
-    vtAPIDomainChannel.send(VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse))
+    withLock L[]:
+          results[].add VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse)
 
 
 proc vtDomainLookup(apiKey: string, domainList: string, jsonOutput: string = "", output: string, rateLimit: int = 4, quiet: bool = false) =
@@ -88,25 +86,22 @@ proc vtDomainLookup(apiKey: string, domainList: string, jsonOutput: string = "",
     headers["x-apikey"] = apiKey
     bar[0].total = len(lines)
     bar.setup()
-    vtAPIDomainChannel.open()
-
-    for domain in lines:
-        inc bar
-        bar.update(1000000000) # refresh every second
-        spawn queryDomainAPI(domain, headers) # run queries in parallel
-
-        # Sleep to respect the rate limit.
-        sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
-
-    for domain in lines:
-        let vtResult: VirusTotalResult = vtAPIDomainChannel.recv() # get results of queries executed in parallel
-        seqOfResultsTables.add(vtResult.resTable)
-        jsonResponses.add(vtResult.resJson)
-        if vtResult.resTable["Response"] == "200" and parseInt(vtResult.resTable["MaliciousCount"]) > 0:
+    var m = createMaster()
+    var results = newSeq[VirusTotalResult]()
+    var L = initTicketLock() # protects `results`
+    m.awaitAll:
+      for domain in lines:
+          inc bar
+          bar.update(1000000000) # refresh every second
+          m.spawn queryDomainAPI(domain, headers, addr results, addr L) # run queries in parallel
+          # Sleep to respect the rate limit.
+          sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
+
+    for r in results:
+        seqOfResultsTables.add(r.resTable)
+        jsonResponses.add(r.resJson)
+        if r.resTable["Response"] == "200" and parseInt(r.resTable["MaliciousCount"]) > 0:
           totalMaliciousDomainCount += 1
-
-    sync()
-    vtAPIDomainChannel.close()
     bar.finish()
 
     echo ""

src/takajopkg/vtHashLookup.nim

@@ -2,9 +2,7 @@
 # https://blog.virustotal.com/2021/08/introducing-known-distributors.html
 # Add output not found to txt file
 
-var vtAPIHashChannel: Channel[VirusTotalResult] # channel for receiving parallel query results
-
-proc queryHashAPI(hash:string, headers: httpheaders.HttpHeaders) {.thread.} =
+proc queryHashAPI(hash:string, headers: httpheaders.HttpHeaders, results: ptr seq[VirusTotalResult]; L: ptr TicketLock) =
     let response = get("https://www.virustotal.com/api/v3/files/" & hash, headers)
     var jsonResponse = %* {}
     var singleResultTable = newTable[string, string]()
@@ -24,8 +22,8 @@ proc queryHashAPI(hash:string, headers: httpheaders.HttpHeaders) {.thread.} =
         singleResultTable["MaliciousCount"] = getJsonValue(jsonResponse, @["data", "attributes", "last_analysis_stats", "malicious"])
         singleResultTable["HarmlessCount"] = getJsonValue(jsonResponse, @["data", "attributes", "last_analysis_stats", "harmless"])
         singleResultTable["SuspiciousCount"] = getJsonValue(jsonResponse, @["data", "attributes", "last_analysis_stats", "suspicious"])
-
-    vtAPIHashChannel.send(VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse))
+    withLock L[]:
+          results[].add VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse)
 
 
 proc vtHashLookup(apiKey: string, hashList: string, jsonOutput: string = "", output: string = "", rateLimit: int = 4, quiet: bool = false) =
@@ -78,25 +76,23 @@ proc vtHashLookup(apiKey: string, hashList: string, jsonOutput: string = "", out
     headers["x-apikey"] = apiKey
     bar[0].total = len(lines)
     bar.setup()
-    vtAPIHashChannel.open()
-
-    for hash in lines:
-        inc bar
-        bar.update(1000000000) # refresh every second
-        spawn queryHashAPI(hash, headers) # run queries in parallel
-
-        # Sleep to respect the rate limit.
-        sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
-
-    for hash in lines:
-        let vtResult: VirusTotalResult = vtAPIHashChannel.recv() # get results of queries executed in parallel
-        seqOfResultsTables.add(vtResult.resTable)
-        jsonResponses.add(vtResult.resJson)
-        if vtResult.resTable["Response"] == "200" and parseInt(vtResult.resTable["MaliciousCount"]) > 0:
+    var m = createMaster()
+    var results = newSeq[VirusTotalResult]()
+    var L = initTicketLock() # protects `results`
+    m.awaitAll:
+      for hash in lines:
+          inc bar
+          bar.update(1000000000) # refresh every second
+          m.spawn queryHashAPI(hash, headers, addr results, addr L) # run queries in parallel
+          # Sleep to respect the rate limit.
+          sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
+
+    for r in results:
+        seqOfResultsTables.add(r.resTable)
+        jsonResponses.add(r.resJson)
+        if r.resTable["Response"] == "200" and parseInt(r.resTable["MaliciousCount"]) > 0:
           totalMaliciousHashCount += 1
 
-    sync()
-    vtAPIHashChannel.close()
     bar.finish()
 
     echo ""
@@ -130,8 +126,8 @@ proc vtHashLookup(apiKey: string, hashList: string, jsonOutput: string = "", out
                 else:
                     outputFile.write(",")
             outputFile.write("\p")
-        let fileSize = getFileSize(output)
         outputFile.close()
+        let fileSize = getFileSize(output)
         echo ""
         echo "Saved CSV results to " & output & " (" & formatFileSize(fileSize) & ")"
 

src/takajopkg/vtIpLookup.nim

@@ -1,8 +1,6 @@
 # TODO: add SAN array info
 
-var vtIpAddressChannel: Channel[VirusTotalResult] # channel for receiving parallel query results
-
-proc queryIpAPI(ipAddress:string, headers: httpheaders.HttpHeaders) {.thread.} =
+proc queryIpAPI(ipAddress:string, headers: httpheaders.HttpHeaders, results: ptr seq[VirusTotalResult]; L: ptr TicketLock) =
     let response = get("https://www.virustotal.com/api/v3/ip_addresses/" & ipAddress, headers)
     var jsonResponse = %* {}
     var singleResultTable = newTable[string, string]()
@@ -37,8 +35,8 @@ proc queryIpAPI(ipAddress:string, headers: httpheaders.HttpHeaders) {.thread.} =
         singleResultTable["SSL-Issuer"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "issuer", "O"])
         singleResultTable["SSL-IssuerCountry"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "issuer", "C"])
         singleResultTable["SSL-CommonName"] = getJsonValue(jsonResponse, @["data", "attributes", "last_https_certificate", "subject", "CN"])
-
-    vtIpAddressChannel.send(VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse,))
+    withLock L[]:
+          results[].add VirusTotalResult(resTable:singleResultTable, resJson:jsonResponse)
 
 
 proc vtIpLookup(apiKey: string, ipList: string, jsonOutput: string = "", output: string, rateLimit: int = 4, quiet: bool = false) =
@@ -91,25 +89,23 @@ proc vtIpLookup(apiKey: string, ipList: string, jsonOutput: string = "", output:
     headers["x-apikey"] = apiKey
     bar[0].total = len(lines)
     bar.setup()
-    vtIpAddressChannel.open()
-
-    for ipAddress in lines:
-        inc bar
-        bar.update(1000000000) # refresh every second
-        spawn queryIpAPI(ipAddress, headers) # run queries in parallel
-
-        # Sleep to respect the rate limit.
-        sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
-
-    for ipAddress in lines:
-        let vtResult: VirusTotalResult = vtIpAddressChannel.recv() # get results of queries executed in parallel
-        seqOfResultsTables.add(vtResult.resTable)
-        jsonResponses.add(vtResult.resJson)
-        if vtResult.resTable["Response"] == "200" and parseInt(vtResult.resTable["MaliciousCount"]) > 0:
+    var m = createMaster()
+    var results = newSeq[VirusTotalResult]()
+    var L = initTicketLock() # protects `results`
+    m.awaitAll:
+      for ipAddress in lines:
+          inc bar
+          bar.update(1000000000) # refresh every second
+          m.spawn queryIpAPI(ipAddress, headers, addr results, addr L) # run queries in parallel
+          # Sleep to respect the rate limit.
+          sleep(int(timePerRequest * 1000)) # Convert to milliseconds.
+
+    for r in results:
+        seqOfResultsTables.add(r.resTable)
+        jsonResponses.add(r.resJson)
+        if r.resTable["Response"] == "200" and parseInt(r.resTable["MaliciousCount"]) > 0:
           totalMaliciousIpAddressCount += 1
 
-    sync()
-    vtIpAddressChannel.close()
     bar.finish()
 
     echo ""

takajo.nimble

@@ -16,4 +16,5 @@ requires "cligen >= 1.5"
 requires "suru#f6f1e607c585b2bc2f71309996643f0555ff6349"
 requires "puppy >= 2.1.0"
 requires "termstyle"
-requires "nancy"
+requires "nancy"
+requires "malebolgia"