Masscan parsing fun

masscan_partser.py

import json
import traceback

class MasscanJsonFile:
	def __init__(self, filename, in_memory = True):
		self.filename = filename
		self.in_memory = in_memory

		self.hosts_port = {}
		self.port_hosts = {}
		self.__merge = []

		if self.in_memory == True:
			self.__parse_file()

	def __radd__(self, x):
		return x + self

	def __add__(self, x):
		if not isinstance(x, MasscanJsonFile):
			raise Exception('Can only add MasscanJsonFile objects!')

		if self.in_memory != x.in_memory:
			raise Exception('Both files needs to be opened in the same "in_memory" value!!')

		if self.in_memory == False:
			self.__merge.append(x)

		else:
			for host in x.hosts_port:
				if host not in self.hosts_port:
					self.hosts_port[host] = []
				self.hosts_port[host] += x.hosts_port[host]

			for port in x.port_hosts:
				if port not in self.port_hosts:
					self.port_hosts[port] = []
				self.port_hosts[port] += x.port_hosts[port]

		return self


	def set_merge(self, mjf):
		"""
		Appends a different MasscanJsonFile for merging.
		WARNING! Both files needs to be opened in the same "in_memory" value!!
		"""
		if self.in_memory != mjf.in_memory:
			raise Exception('Both files needs to be opened in the same "in_memory" value!!')

		if self.in_memory == False:
			self.__merge.append(mjf)
		else:
			for host in mjf.hosts_port:
				if host not in self.hosts_port:
					self.hosts_port[host] = []
				self.hosts_port[host] += mjf.hosts_port[host]

			for port in mjf.port_hosts:
				if port not in self.port_hosts:
					self.port_hosts[port] = []
				self.port_hosts += mjf.port_hosts[port]

	def yield_data(self):
		"""
		Reads the result for the file, one result at a time, and yileds the ip-port tuple.
		"""
		with open(self.filename,'r') as f:
			cnt = 0
			for line in f:
				cnt += 1
				if line[0] != '{' or line[0] == ',':
					continue
				try:
					data = json.loads(line)
				except Exception as e:
					print('Error in file %s on line %d! Skipping line...' % (self.filename, cnt))
					#traceback.print_exc()
					continue
				else:
					yield (data['ip'], data['ports'][0]['port'])

		for mjf in self.__merge:
			for ip, port in mjf.yield_data():
				yield ip, port
				

	def __parse_file(self):
		"""
		If in-memory is set to True, this function will parse the whole json file 
		and create two dictionaries, one for storing host-port and another for storing port-hosts lists
		"""
		with open(self.filename, 'r') as f:
			data = json.load(f)

		for el in data:
			if el['ip'] not in self.hosts_port:
				self.hosts_port[el['ip']] = []

			for pe in el['ports']:
				if pe['status'] == 'open':
					if pe['port'] not in self.port_hosts:
						self.port_hosts[pe['port']] = []
					self.port_hosts[pe['port']].append(el['ip'])
					self.hosts_port[el['ip']].append(pe['port'])

	def get_hosts_for_port(self, portno):
		"""
		Yields hosts(ip) for given port open
		"""
		if self.in_memory == True:
			if portno in self.port_hosts:
				for ip in self.port_hosts[portno]:
					yield ip

		else:
			for host, port in self.yield_data():
				if port == portno:
					yield host


	def get_ports_for_host(self, ip):
		"""
		Yields ports open for given host(ip)
		"""
		if self.in_memory == True:
			if ip in self.hosts_port:
				for portno in self.hosts_port[ip]:
					yield portno

		else:
			for host, port in self.yield_data():
				if host == ip:
					yield port

		for mjf in self.__merge:
			yield mjf.get_ports_for_host(ip)

	def write_ports_for_host(self, filename, host):
		"""
		Creates one file, and writes all ports open on a given host. 
		One port per line.
		"""
		if filename:
			with open(filename, 'w', newline = '') as f:
				for port in self.get_ports_for_host(host):
					f.write(port + '\r\n')
		else:
			for port in self.get_ports_for_host(host):
				print(port)

	def write_all_for_host(self, base_filename):
		"""
		Creates mutiple files, one per host. Writes all ports open for the given host in the specific file.
		"""
		if self.in_memory == True:
			if base_filename:
				for host in self.hosts_port:
					filename = '%s_%d.txt' % (base_filename, host)
					self.write_ports_for_host(filename, host)
			else:
				for host in self.hosts_port:
					print('====== HOST %s ======' % portno)
					self.write_ports_for_host(None, host)

		else:
			if base_filename:
				for host, port in self.yield_data():
					filename = '%s_%d.txt' % (base_filename, host)
					with open(filename, 'a', newline = '') as f:
						f.write(port + '\r\n')
			else:
				raise Exception('Cant print to STDOUT in this case. Please specify a base filename to write to!')


	def write_hosts_for_port(self, filename, portno):
		"""
		Creates one file, and writes all hosts where the given port is open. 
		One host per line.
		"""
		if filename:
			with open(filename, 'w', newline = '') as f:
				for ip in self.get_hosts_for_port(portno):
					f.write(ip + '\r\n')
		else:
			for ip in self.get_hosts_for_port(portno):
				print(ip)

	def write_all_for_ports(self, base_filename):
		"""
		Creates multiple files, one per port. 
		Each file will have all host addresses where the specific port is open.
		One host per line.
		"""
		if self.in_memory == True:
			if base_filename:
				for portno in self.port_hosts:
					filename = '%s_%d.txt' % (base_filename, portno)
					self.write_hosts_for_port(filename, portno)
			else:
				for portno in self.port_hosts:
					print('====== PORT %d ======' % portno)
					self.write_hosts_for_port(None, portno)

		else:
			if base_filename:
				for host, port in self.yield_data():
					filename = '%s_%d.txt' % (base_filename, port)
					with open(filename, 'a', newline = '') as f:
						f.write(host + '\r\n')
			else:
				raise Exception('Cant print to STDOUT in this case. Please specify a base filename to write to!')

	def write_all(self, filename = None):
		"""
		Writes a single file with all open port in "host:port" format per line 
		"""
		if self.in_memory == True:
			if filename:
				with open(filename, 'w', newline = '') as f:
					for portno in self.port_hosts:
						for ip in self.port_hosts[portno]:
							f.write('%s:%s\r\n' % (ip, portno))
			else:
				for portno in self.port_hosts:
					for ip in self.port_hosts[portno]:
						print('%s:%s' % (ip, portno))
		else:
			if filename:
				with open(filename, 'w', newline = '') as f:
					for host, port in self.yield_data():
						f.write('%s:%s\r\n' % (host, port))
			else:
				for host, port in self.yield_data():
					print('%s:%s' % (host, port))

def run():
	import argparse
	import glob

	parser = argparse.ArgumentParser(description='Tool to process Masscan JSON file format')
	parser.add_argument('-v', '--verbose', action='count', default=0, help='Increase verbosity, can be stacked. Doesnt do anyhting.')
	parser.add_argument('-d', '--disable-in-memory-parsing', action='store_false', help='Disable in-memory parsing. Use this for large AND/OR broken files.')
	parser.add_argument('-f','--in-file',  action='append', help='Input JSON file. Can be stacked.')
	parser.add_argument('-g','--glob', help='Pattern to search for files to parse for. Example: "*.json" will parse all json files in the current directory')


	subparsers = parser.add_subparsers(help = 'commands')
	subparsers.required = True
	subparsers.dest = 'command'

	all_group = subparsers.add_parser('all', help='Generate a list of all ports on all hosts. Output firmat: <ip>:<port>')
	
	host_group = subparsers.add_parser('host', help='List open ports on given host')
	host_group.add_argument('host', help='Host IP address to filter for. "ALL" will create one txt file per host and list all open ports there!')

	port_group = subparsers.add_parser('port', help='List all hosts that have given port open')
	port_group.add_argument('port', help='Port to filter for. "ALL" will create one txt file per port and list all hosts there that have specific port open')

	parser.add_argument('-o','--out-file', help='Output file name. If not specified will print to stdout, or fail if "disable-in-memory-parsing" is set and "ALL" filter is used')

	args = parser.parse_args()

	files = None

	if args.glob:
		for filename in glob.glob(args.glob):
			print('[+] Parsing file %s' % filename)
			if not files:
				files = MasscanJsonFile(filename, args.disable_in_memory_parsing)
			else:
				files += MasscanJsonFile(filename, args.disable_in_memory_parsing)

	if args.in_file:
		for filename in args.in_file:
			print('[+] Parsing file %s' % filename)
			if not files:
				files = MasscanJsonFile(filename, args.disable_in_memory_parsing)
			else:
				files += MasscanJsonFile(filename, args.disable_in_memory_parsing)

	if args.command == 'all':
		files.write_all(args.out_file)

	if args.command == 'host':
		if args.host.upper() == 'ALL':
			files.write_all_for_host(args.out_file)
		else:
			files.write_ports_for_host(args.out_file, args.host)

	if args.command == 'port':
		if args.port.upper() == 'ALL':
			files.write_all_for_ports(args.out_file)
		else:
			files.write_hosts_for_port(args.out_file, args.port)
      
if __name__ == '__main__':
	run()