PREPARE/prepare/scanner_wrapper.py

import os
import logging
import collections
import threading
import subprocess
import time

from subprocess import PIPE, TimeoutExpired

logger = logging.getLogger("pra_framework")

class ZmapWrapper(object):
	def __init__(self, **kwargs):
		# there ary many arguments so we take a more efficient approach
		timestamp = time.strftime("%Y_%M_%d__%H_%M_%S")

		# Default parameters to be set. This is more efficient flexible
		# than explicitly setting them in __init__() via self._xyz...
		prop_defaults = {
			"port_target": 1234,
			"filename_output_csv" : None,
			"filename_blacklist_target_ip": None,
			# TODO: just for online scans
			"filename_whitelist_target_ip": None,
			"dirname_logs": None,
			# scan options
			"rate_kbit_per_s": 1000,
			"probes_per_target": 1,
			"timeout_wait": 0,
			# network options
			"interface_name": "eth0",
			"mac_source": None,
			"mac_gw": None,
			# probe module, we only use the pra module
			"probe_module": "tcp_synscan_pra",
			#"probe_module": "tcp_synscan",
			# additional options
			"marker_encoding": 0,
			"markervalue": None,
			"markerbits_value": 32,
			"markerbits_checksum": 0,
			"markerbits_dst_small" : 0,
			# disable status outputs, directly exit scanner after finishing (for small address ranges to speed up)
			"disable_monitor": 0,
			"filename_config": None,
			# target adresses as string: "1.2.3.4 1.2.4.0/24 ..."
			"target_addresses": None,
			"verbosity": 3,
			#
			"ready_callback": None,
			"name_timestamp": timestamp,
			# directory which contains "src/zmap wihtout trailing slash"
			"base_dir_zmap": "../zmap",
			# avoid establishing process communication and don't wait for process to finish
			"fast_mode": False,
			"scan_timeout": -1
		}

		for (prop, default) in prop_defaults.items():
			setattr(self, prop, kwargs.get(prop, default))

		if self.target_addresses is None:
			raise Exception("Not address given to scan")

		self.filename_zmap_bin = os.path.join(self.base_dir_zmap, "src/zmap")
		# use absolute path
		#self.filename_zmap_bin = os.path.abspath(self.filename_zmap_bin)
		# logger.debug("zmap binary: %s" % self.filename_zmap_bin)

		"""
		if self.filename_blacklist_target_ip is None:
			# TODO: change to full blacklist file
			self.filename_zmap_blacklist = os.path.join(self.base_dir_zmap, "blacklist.conf")
			#self.filename_blacklist_target_ip = os.path.join(self.base_dir_zmap, "blacklist.conf_minimal")
		"""
		self._scan_process = None
		self._is_running = False

	def start(self):
		"""
		Start a port scan. This call is blocking until the
		process is finished. Call stop() to terminate process earlier.
		"""
		if self._is_running:
			logger.warning("scanner is allready running")
			return
		# aggregate scan parameters
		cmd = [self.filename_zmap_bin]

		# basic arguments
		# this is generally not needed when using destination port marker but we set is per default
		cmd.extend(["-p", "%d" % self.port_target])
		if self.filename_blacklist_target_ip is not None:
			cmd.extend(["-b", self.filename_blacklist_target_ip])
		"""
		if self.filename_whitelist_target_ip is not None:
			cmd.extend(["-w", self.filename_whitelist_target_ip])
		"""
		if self.filename_output_csv is not None:
			cmd.extend(["-o", self.filename_output_csv])
			# data output
			cmd.extend(["-O", "csv"])
			cmd.extend(["-f", "classification,saddr,daddr,daddr_inner_icmp,sport,dport,success"])
		# scan options
		cmd.extend(["-B", "%dK" % self.rate_kbit_per_s])
		cmd.extend(["-P", "%d" % self.probes_per_target])
		if self.filename_output_csv is None:
			cmd.extend(["-c", "%d" % self.timeout_wait])
		else:
			# wait at minimum 2 seconds at end of scan to receive feedback
			cmd.extend(["-c", "%d" % 2])
		#cmd.extend(["--hello-timeout", self._timeout_syn])
		# network options
		if self.mac_gw is not None:
			cmd.extend(["-G", self.mac_gw])
		if self.mac_source is not None:
			cmd.extend(["--source-mac", self.mac_source])
		cmd.extend(["-i", self.interface_name])
		# probe modules
		cmd.extend(["-M", self.probe_module])
		# additional options
		cmd.extend(["--marker_encoding", "%d" % self.marker_encoding])
		if self.markervalue is not None:
			cmd.extend(["--use_markervalue", "1"])
			cmd.extend(["--markervalue", "%d" % self.markervalue])
		cmd.extend(["--markerbits_value", "%d" % self.markerbits_value])
		cmd.extend(["--markerbits_checksum", "%d" % self.markerbits_checksum])
		cmd.extend(["--disable_monitor", "%d" % self.disable_monitor])

		if self.filename_config is not None:
			cmd.extend(["-C", self.filename_config])
		cmd.extend(["-v", "%d" % self.verbosity])
		# TODO: activate for zmap internal Logging
		if self.dirname_logs is not None:
			cmd.extend(["-L", self.dirname_logs])
		# limit scan time
		if self.scan_timeout > 0:
			cmd.extend(["-t", "%d" % self.scan_timeout])
		# two cores are enough to achieve a maximum speed
		#cmd.extend(["--cores", "0"])
		#cmd.extend(["--cores", "0,1"])
		cmd.extend(["--cores", "0,1,2"])
		# addresses as list
		cmd.extend(self.target_addresses)

		logger.debug("starting scanner: " + " ".join(cmd))

		self._is_running = True

		#return

		# start the process
		try:
			#self._scan_process = subprocess.Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
			if not self.fast_mode:
				self._scan_process = subprocess.Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
			else:
				self._scan_process = subprocess.Popen(cmd)
		except Exception:
			logger.warning("could not start process, killing it")
			self._scan_process.kill()

		#if self.fast_mode:
		#	return
		# if self._scan_process.returncode != 0:
		#	logger.warning("something went wrong while scanning, check output")
		# wait until scan has finished
		try:
			self._scan_process.wait()
			self._check_process_output()
		except TimeoutExpired:
			# this will update date for status retrieval (time left)
			self._check_process_output()

		if self._is_running:
			self.stop()

		if self.ready_callback is not None:
			# self._extract_ips()
			self.ready_callback(self)

	def _check_process_output(self):
		try:
			outs, errs = self._scan_process.communicate(timeout=1)
			#if self.disable_monitor == 0:
			#	logger.debug("standard output\n:" + outs.decode("ASCII"))
			#	logger.debug("error output:\n" + errs.decode("ASCII"))
		except:
			pass

	def stop(self):
		if not self._is_running:
			logger.warning("scanner was not running (anymore?)")
			return
		self._is_running = False
		# avoid ResourceWarning: unclosed file <_io.BufferedReader name=7>
		try:
			self._scan_process.stdout.close()
			self._scan_process.stdin.close()
			self._scan_process.stderr.close()
		except:
			pass

		try:
			# self._scan_process.terminate()
			# do it the hard way, we don't need any resume storing
			self._scan_process.kill()
		except ProcessLookupError as e:
			# logger.warning("could not terminate process")
			# logger.warning(e)
			# assume process already finished
			pass