diff --git a/src/takajo.nim b/src/takajo.nim index f9e0238d..1947553f 100644 --- a/src/takajo.nim +++ b/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 diff --git a/src/takajopkg/vtDomainLookup.nim b/src/takajopkg/vtDomainLookup.nim index 5d8bbd21..d7e36941 100644 --- a/src/takajopkg/vtDomainLookup.nim +++ b/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 "" diff --git a/src/takajopkg/vtHashLookup.nim b/src/takajopkg/vtHashLookup.nim index ba9efb9f..74d4d3d0 100644 --- a/src/takajopkg/vtHashLookup.nim +++ b/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) & ")" diff --git a/src/takajopkg/vtIpLookup.nim b/src/takajopkg/vtIpLookup.nim index 0fb391a0..b5b52e47 100644 --- a/src/takajopkg/vtIpLookup.nim +++ b/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 "" diff --git a/takajo.nimble b/takajo.nimble index 6074623d..a4a4346d 100644 --- a/takajo.nimble +++ b/takajo.nimble @@ -16,4 +16,5 @@ requires "cligen >= 1.5" requires "suru#f6f1e607c585b2bc2f71309996643f0555ff6349" requires "puppy >= 2.1.0" requires "termstyle" -requires "nancy" \ No newline at end of file +requires "nancy" +requires "malebolgia" \ No newline at end of file