Add new modules and optimize DB interaction (v1.0.0)

.gitignore

@@ -26,4 +26,4 @@ venv.bak/
 # Custom files
 !/das/*
 !/.github/*
-TODO.txt
+*.txt

README.md

@@ -7,7 +7,7 @@
 </p>
 
 <p align="center">
-  <a href="https://github.com/snovvcrash/DivideAndScan/blob/main/pyproject.toml#L3"><img src="https://img.shields.io/badge/version-0.3.8-success" alt="version" /></a>
+  <a href="https://github.com/snovvcrash/DivideAndScan/blob/main/pyproject.toml#L3"><img src="https://img.shields.io/badge/version-1.0.0-success" alt="version" /></a>
   <a href="https://github.com/snovvcrash/DivideAndScan/search?l=python"><img src="https://img.shields.io/badge/python-3.9-blue?logo=python&logoColor=white" alt="python" /></a>
   <a href="https://www.codacy.com/gh/snovvcrash/DivideAndScan/dashboard?utm_source=github.com&amp;utm_medium=referral&amp;utm_content=snovvcrash/DivideAndScan&amp;utm_campaign=Badge_Grade"><img src="https://app.codacy.com/project/badge/Grade/35f0bdfece9846d7aab3888b01642813" alt="codacy" /></a>
   <a href="https://github.com/snovvcrash/DivideAndScan/actions/workflows/publish-to-pypi.yml"><img src="https://github.com/snovvcrash/DivideAndScan/actions/workflows/publish-to-pypi.yml/badge.svg" alt="pypi" /></a>
@@ -104,6 +104,14 @@ sudo eget -s linux/amd64 v-byte-cpu/sx --to /opt/sx
 sudo ln -sv /opt/sx/sx /usr/local/bin/sx
 ```
 
+#### dnsx
+
+```bash
+sudo mkdir /opt/pd
+sudo eget -s linux/amd64 projectdiscovery/dnsx --to /opt/pd
+sudo ln -sv /opt/pd/dnsx /usr/local/bin/dnsx
+```
+
 ### Installation
 
 DivideAndScan is available on PyPI as `divideandscan`, though I recommend installing it from GitHub with [pipx](https://github.com/pipxproject/pipx) in order to always have the bleeding-edge version:
@@ -320,7 +328,7 @@ class AddPortscanOutput(IAddPortscanOutput):
 ## Help
 
 ```
-usage: das [-h] [-db DB] {add,scan,dns,report,parse,draw,tree,help} ...
+usage: das [-h] [-db DB] {db,add,scan,dns,report,parse,draw,tree,help} ...
 
  -----------------------------------------------------------------------------------------------
 |  ________  .__      .__    .___        _____              .____________                       |
@@ -333,10 +341,11 @@ usage: das [-h] [-db DB] {add,scan,dns,report,parse,draw,tree,help} ...
  -----------------------------------------------------------------------------------------------
 
 positional arguments:
-  {add,scan,dns,report,parse,draw,tree,help}
+  {db,add,scan,dns,report,parse,draw,tree,help}
+    db                  utilities for manual DB manipulations
     add                 run a full port scan and add the output to DB
     scan                run targeted Nmap scans against hosts and ports from DB
-    dns                 resolve "A" domain names into IP addresses and update DB items with them
+    dns                 map domain names from an input file to corresponding IP addresses from the DB
     report              merge separate Nmap outputs into a single report (https://github.com/CBHue/nMap_Merger)
     parse               parse raw Nmap XML reports by service names and print entries in format {service}://{host}:{port}}
     draw                visualize Nmap XML reports (https://github.com/jor6PS/DrawNmap)

das/common.py

@@ -2,9 +2,11 @@
 
 __author__ = '@snovvcrash'
 __site__ = 'https://github.com/snovvcrash/DivideAndScan'
-__version__ = '0.3.8'
+__version__ = '1.0.0'
 
 import time
+import shlex
+import subprocess
 from datetime import datetime, timedelta
 
 BANNER = """\
@@ -108,3 +110,18 @@ def print_separator(msg, prefix):
 		:type prefix: str
 		"""
 		print(f'\033[0;31m{SEP} \033[0;32m({prefix}) \033[1;32m{msg}\033[0;31m {SEP}\033[0m')
+
+
+def run_command(command):
+    process = subprocess.Popen(shlex.split(command), stdout=subprocess.PIPE)
+
+    result = ''
+    while True:
+        stdout = process.stdout.readline().decode()
+        if stdout == '' and process.poll() is not None:
+            break
+        if stdout:
+            result += stdout
+
+    #result = process.poll()
+    return result

das/divideandscan.py

@@ -8,6 +8,7 @@
 from argparse import ArgumentParser, RawTextHelpFormatter, RawDescriptionHelpFormatter
 
 import das.common
+from das.db import DB
 from das.scan import ScanShow, ScanRun
 from das.dns import DNS
 from das.report import NmapMerger
@@ -27,6 +28,16 @@ def parse_args():
 
 	subparser = parser.add_subparsers(dest='subparser')
 
+	db_epilog = """
+	examples:
+	  das -db testdb db scan.txt [-d domains.txt] [--create]
+	""".replace('\t', '')
+	db_parser = subparser.add_parser('db', formatter_class=RawDescriptionHelpFormatter, epilog=db_epilog, help='utilities for manual DB manipulations')
+	db_parser.add_argument('scan_output', action='store', type=str, help='path to a newline-separated text file with scan output to fill the DB with')
+	db_parser.add_argument('-d', '--domains', action='store', type=str, help='path to a newline-separated text file with with domain names to fill the DB with')
+	group_action = db_parser.add_mutually_exclusive_group(required=True)
+	group_action.add_argument('--create', action='store_true', default=False, help='create TinyDB from a generic scan output and (optionally) a list of domain names')
+
 	add_epilog = """
 	examples:
 	  das add nmap '-v -n -Pn -e eth0 --min-rate 1000 -T4 -iL hosts.txt -p1-49151 --open'
@@ -70,10 +81,13 @@ def parse_args():
 
 	dns_epilog = """
 	examples:
-	  das dns domains.txt
+	  das dns domains.txt [--resolve|--update]
 	""".replace('\t', '')
-	dns_parser = subparser.add_parser('dns', formatter_class=RawDescriptionHelpFormatter, epilog=dns_epilog, help='resolve "A" domain names into IP addresses and update DB items with them')
+	dns_parser = subparser.add_parser('dns', formatter_class=RawDescriptionHelpFormatter, epilog=dns_epilog, help='map domain names from an input file to corresponding IP addresses from the DB')
 	dns_parser.add_argument('domains', action='store', type=str, help='path to a newline-separated text file with domain names to resolve')
+	group_action = dns_parser.add_mutually_exclusive_group(required=True)
+	group_action.add_argument('--resolve', action='store_true', default=False, help='resolve "A" domain names into IP addresses and update DB items with them')
+	group_action.add_argument('--update', action='store_true', default=False, help='update existing DB with new domains names from an input file')
 
 	report_epilog = """
 	examples:
@@ -127,7 +141,7 @@ def main():
 	args = parse_args()
 
 	if len(sys.argv) == 1:
-		print('usage: __main__.py [-h] {add,scan,dns,report,parse,draw,tree,help} ...\n')
+		print('usage: __main__.py [-h] {db,add,scan,dns,report,parse,draw,tree,help} ...\n')
 		print(BANNER)
 		sys.exit(0)
 
@@ -136,7 +150,17 @@ def main():
 	if (args.subparser == 'add' and not Path(args.scanner_args).is_file) or args.subparser == 'scan' and not args.show:
 		logger.start_timer()
 
-	if args.subparser == 'add':
+	if args.subparser == 'db':
+		(Path.home() / '.das' / 'db' / 'raw').mkdir(parents=True, exist_ok=True)
+		P = Path.home() / '.das' / 'db' / f'{args.db}.json'
+
+		db = DB(str(P))
+		if args.create:
+			logger.print_info(f'Creating DB -> {P.resolve()}')
+			num_created = db.create_generic(Path(args.scan_output), Path(args.domains))
+			logger.print_success(f'Successfully created DB with {num_created} hosts')
+
+	elif args.subparser == 'add':
 		(Path.home() / '.das' / 'db' / 'raw').mkdir(parents=True, exist_ok=True)
 
 		module_name = Path(args.scanner_name).name
@@ -152,7 +176,7 @@ def main():
 		P = Path.home() / '.das' / 'db' / f'{args.db}.json'
 
 		apo = AddPortscanOutput(str(P), args.rm, args.scanner_name, args.scanner_args)
-		portscan_out, num_of_hosts = apo.parse()
+		portscan_out, num_hosts = apo.parse()
 
 		if P.exists():
 			logger.print_info(f'Using DB -> {P.resolve()}')
@@ -161,7 +185,7 @@ def main():
 		if P.exists():
 			logger.print_info(f'Raw port scanner output -> {P.resolve()}')
 
-		logger.print_success(f'Successfully updated DB with {num_of_hosts} hosts')
+		logger.print_success(f'Successfully updated DB with {num_hosts} hosts')
 
 	elif args.subparser == 'scan':
 		if not args.show:
@@ -198,8 +222,12 @@ def main():
 		if P.exists():
 			logger.print_info(f'Using DB -> {P.resolve()}')
 
-		d = DNS(str(P), Path(args.domains))
-		d.resolve()
+		dns = DNS(str(P), Path(args.domains))
+		if args.resolve:
+			dns.resolve()
+		elif args.update:
+			num_updated = dns.update()
+			logger.print_success(f'Successfully updated {num_updated} hosts from DB with domain names')
 
 	elif args.subparser == 'report':
 		output = {'oA': args.oA, 'oX': args.oX, 'oN': args.oN, 'oG': args.oG}

das/dns.py

@@ -1,20 +1,26 @@
 #!/usr/bin/env python3
 
+import os
+import tempfile
+from pathlib import Path
+from collections import defaultdict
+
 import dns.resolver
 from tinydb import TinyDB, Query
 from tinydb.operations import add
 
-from das.common import Logger
+from das.common import Logger, run_command
+
 
 class DNS:
-	"""Class for associating domain names (taken from an input file) with corresponding IP addresses (taken from the DB)."""
+	"""Class for mapping domain names (taken from an input file) to corresponding IP addresses (taken from the DB)."""
 
 	def __init__(self, db_path, domains_path):
 		"""
 		Constructor.
 
 		:param db_path: a TinyDB database file path
-		:type db_path: pathlib.PosixPath
+		:type db_path: str
 		:param domains_path: an input file path with newline-separated domain names
 		:type domains_path: pathlib.PosixPath
 		"""
@@ -51,3 +57,43 @@ def resolve(self):
 				self.db.update(add('domains', [domain]), doc_ids=doc_ids)
 			elif not found:
 				Logger.print_warning(f'None -> {domain}')
+
+	def update(self):
+		"""Update existing DB with new domains names (taken from an input file).
+
+		:return: number of hosts updated with domain names in the DB
+		:rtype: int
+		"""
+
+		with tempfile.NamedTemporaryFile('w', suffix='.txt') as tmp:
+			dnsx_path = Path(tempfile.gettempdir()) / 'dnsx.txt'
+			domains_punycode = [domain.encode('idna').decode() + '\n' for domain in self.domains]
+			tmp.writelines(domains_punycode)
+			run_command(f'dnsx -l {tmp.name} -re -silent -o {dnsx_path}')
+
+		with open(dnsx_path) as f:
+			dnsx = f.read().splitlines()
+
+		Logger.print_info(f'Resolved {len(dnsx)} DNS records')
+		os.remove(dnsx_path)
+
+		domains = defaultdict(set)
+		for line in dnsx:
+			domain, ip = line.split()
+			domain = domain.encode().decode('idna')
+			ip = ip.replace('[', '').replace(']', '')
+			domains[ip].add(domain)
+
+		items, hosts = self.db.all(), set()
+		for i in range(len(items)):
+			ip = items[i]['ip']
+			if ip in domains:
+				t = set(items[i]['domains'])
+				t.update(domains[ip])
+				items[i]['domains'] = list(t)
+				hosts.add(ip)
+
+		self.db.truncate()
+		self.db.insert_multiple(items)
+
+		return len(hosts)

das/parsenmap.py

@@ -8,6 +8,7 @@
 
 from das.common import Logger
 
+
 class NmapParser:
 	"""Class for parsing Nmap XML reports by service names and print entries in format {service}://{host}:{port}}."""
 
@@ -27,8 +28,8 @@ def __init__(self, db_path, services, dns, raw_output=False):
 		:rtype: das.report.NmapParser
 		"""
 		self.services = services.split(',')
-
-		self.db = None
+		self.dns = dns
+
 		if dns:
 			self.db = TinyDB(db_path)
 			self.ip_domains_dict = {}
@@ -58,7 +59,7 @@ def parse(self):
 						if nm[ip]['tcp'][port]['state'] == 'open':
 							service = nm[ip]['tcp'][port]['name']
 							if service in self.services:
-								if self.db:
+								if self.dns:
 									domains = self.ip_domains_dict[ip]
 									if domains:
 										for domain in domains:

das/parsers/__init__.py

@@ -19,7 +19,7 @@ def __init__(self, db_path, rm, scanner_name, scanner_args):
 		Constructor.
 
 		:param db_path: a TinyDB database file path
-		:type db_path: pathlib.PosixPath
+		:type db_path: str
 		:param rm: a flag showing if we need to drop the DB before updating its values
 		:type rm: bool
 		:param scanner_name: name of the port scanner to run

das/parsers/masscan.py

@@ -11,7 +11,7 @@ def parse(self):
 		:return: a pair of values (portscan raw output filename, number of hosts added to DB)
 		:rtype: tuple
 		"""
-		hosts = set()
+		items, hosts = [], set()
 		for line in self.portscan_raw:
 			try:
 				ip = line.split()[-1]
@@ -20,10 +20,9 @@ def parse(self):
 				pass
 			else:
 				if proto == 'tcp':
-					item = {'ip': ip, 'port': int(port), 'domains': []}
-					if item not in self.db:
-						self.db.insert(item)
+					items.append({'ip': ip, 'port': int(port), 'domains': []})
+					hosts.add(ip)
 
-				hosts.add(ip)
+		self.db.insert_multiple(items)
 
 		return (self.portscan_out, len(hosts))

das/parsers/naabu.py

@@ -11,17 +11,16 @@ def parse(self):
 		:return: a pair of values (portscan raw output filename, number of hosts added to DB)
 		:rtype: tuple
 		"""
-		hosts = set()
+		items, hosts = [], set()
 		for line in self.portscan_raw:
 			try:
 				ip, port = line.split(':')
 			except Exception:
 				pass
 			else:
-				item = {'ip': ip, 'port': int(port), 'domains': []}
-				if item not in self.db:
-					self.db.insert(item)
-
+				items.append({'ip': ip, 'port': int(port), 'domains': []})
 				hosts.add(ip)
 
+		self.db.insert_multiple(items)
+
 		return (self.portscan_out, len(hosts))

das/parsers/nimscan.py

@@ -11,7 +11,7 @@ def parse(self):
 		:return: a pair of values (portscan raw output filename, number of hosts added to DB)
 		:rtype: tuple
 		"""
-		hosts = set()
+		items, hosts = [], set()
 		for line in self.portscan_raw:
 			try:
 				ip, port = line.split('==>')[1].split()[0].split(':')
@@ -23,10 +23,9 @@ def parse(self):
 			except Exception:
 				pass
 			else:
-				item = {'ip': ip, 'port': int(port), 'domains': []}
-				if item not in self.db:
-					self.db.insert(item)
-
+				items.append({'ip': ip, 'port': int(port), 'domains': []})
 				hosts.add(ip)
 
+		self.db.insert_multiple(items)
+
 		return (self.portscan_out, len(hosts))

das/parsers/nmap.py

@@ -11,18 +11,17 @@ def parse(self):
 		:return: a pair of values (portscan raw output filename, number of hosts added to DB)
 		:rtype: tuple
 		"""
-		hosts = set()
+		items, hosts = [], set()
 		for line in self.portscan_raw:
 			try:
 				ip = line.split()[-1]
 				port, _ = line.split()[3].split('/')  # port, proto
 			except Exception:
 				pass
 			else:
-				item = {'ip': ip, 'port': int(port), 'domains': []}
-				if item not in self.db:
-					self.db.insert(item)
-
+				items.append({'ip': ip, 'port': int(port), 'domains': []})
 				hosts.add(ip)
 
+		self.db.insert_multiple(items)
+
 		return (self.portscan_out, len(hosts))