def do_auth(self):
"""
Handles the agent's authentication.
"""
log.info(f"{self.client_address} initialized Authentication")
challenge = os.urandom(128)
self.request.sendall(Auth(challenge).pack())
# wait for the client response
self.msg = Structure.create(self.request)
if not self.msg:
# something's wrong with the message!
self.send_status(Status.FAILED)
return
hmac_hash = hmac.new(self.server.secret_key, challenge, 'sha512')
digest = hmac_hash.hexdigest().encode("utf-8")
if hmac.compare_digest(digest, self.msg.data):
log.info("Authenticated Successfully")
self.authenticated = True
self.send_status(Status.SUCCESS)
else:
self.send_status(Status.UNAUTHORIZED)
self.request.close()
def dispatcher(self):
"""
Command dispatcher all logic to decode and dispatch the call.
"""
self.msg = Structure.create(self.request)
if not self.msg:
self.connected = False
log.info("Disconnected!")
# mark any running task as interrupted
# so that other agent can take it later
self.ctx.interrupted(self.agent)
return
command_name = f"do_{self.msg.op_code.name.lower()}"
if not hasattr(self, command_name):
self.send_status(Status.FAILED) # invalid command
return
command = getattr(self, command_name)
# the only command authorized for unauthenticated agents
if command_name != "do_AUTH" and not self.authenticated:
self.send_status(Status.UNAUTHORIZED)
self.connected = False
self.request.close()
return
# call the command !
command()
def __getstate__(self):
# Copy the object's state from self.__dict__ which contains
# all our instance attributes. Always use the dict.copy()
# method to avoid modifying the original state.
log.info(f"saving file: {self._path} loc:{self.loc} state")
state = self.__dict__.copy()
# Remove the unpickable entries.
del state['_fd']
return state
def save_context(self, ctx):
"""
Serializes the context to resume later.
:param ctx: instance of `Context`
:type ctx: `Context`
"""
log.info(f"Saving the current context {self.resume_path}")
with open(self.resume_path, 'wb') as rfile:
pickle.dump(ctx, rfile)
def send_status(self, code):
"""
Sends a status code to the server.
:param code:
:type code: `dscan.models.structures.Status`
"""
log.info(f"Sending status code {code}")
response = struct.pack("<B", code)
self.request.sendall(response)
def pop(self, agent):
"""
Gets the next `Task` from the current Active Stage, if their are no
pending `Tasks` to be executed.
Pending tasks are tasks that are canceled or restored from a previous
interrupted session.
If a stage is finished (no more targets), the next stage will take
another stage from the list until its finished.
:param agent:
str with ipaddress and port in ip:port format, this allows the
server to manage multiple agents in one host.
to run multiple clients at once.
:return: A target to scan! `task`
:rtype: `tuple`
"""
with self._lock:
task = None
if agent in self.active:
# This exists to make shore we don't lose targets.
# this would be better if we knew how many tries a target
# had faced. TODO
log.info(f"Agent {agent} is requesting a new task with a "
f"task in execution sending it again!")
task = self.active.get(agent)
task.update(STATUS.SCHEDULED)
return task.as_tuple()[2:]
if len(self.pending) > 0:
task = self.pending.pop(0)
else:
cstage = self.__cstage()
if cstage:
task = cstage.next_task()
if not task:
# the only stage that needs to be finished
# to proceed is
# discovery as the other stages need the
# list of live hosts
if cstage.name != "discovery" or cstage.isfinished:
if cstage.isfinished:
cstage.process_results()
cstage.close()
cstage = self.__cstage(True)
if cstage:
task = cstage.next_task()
# if we have a valid task save it in the active collection
if task:
self.active.update({agent: task})
# the consumers only need scan related information...
return task.as_tuple()[2:]
def run(self, target, options, callback):
"""
Executes the scan on a given target.
:param target:
:param options:
:param callback: callback function to report status to the server.
:return: report object
:rtype: `dscan.models.structures.Report`
"""
self.ctarget = (target, options)
nmap_proc = None
try:
options = " ".join([options, f"-oN {self.report_name('nmap')}"])
nmap_proc = NmapProcess(targets=target,
options=options,
safe_mode=False,
event_callback=self.show_status)
log.info("Nmap scan started Sending success status")
callback(Status.SUCCESS)
rc = nmap_proc.run()
if rc == 0:
# after finished encode and hash the contents for transfer.
self.__inc()
data = nmap_proc.stdout.encode("utf-8")
report_file = self.report_name("xml")
with open(self.report_name("xml"), "wb") as rfile:
rfile.write(data)
rfile.flush()
digest = hashlib.sha512(data).hexdigest()
report = Report(len(data), os.path.basename(report_file),
digest)
self.print(target, 100)
return report
else:
callback(Status.FAILED)
log.error(f"Nmap Scan failed {nmap_proc.stderr}")
except Exception as ex:
log.error(f"something went wrong {ex}")
callback(Status.FAILED)
finally:
if nmap_proc:
nmap_proc.stop()
# orthodox fix NmapProcess is leaving subprocess streams open.
subproc = getattr(nmap_proc, "_NmapProcess__nmap_proc")
if subproc:
subproc.stdout.close()
subproc.stderr.close()
def __setstate__(self, state):
# Restore instance attributes (i.e., _path and nlines ...).
self.__dict__.update(state)
fd = None
# if the loc is 0 then we have an uninitialized stage
# that depends on unfinished stage, will be opened when the first
# time a target is pulled.
if self.loc != 0 and self.exists():
# Restore the previously opened file's state.
log.info(f"restoring file: {self._path} loc:{self.loc} state")
fd = open(self._path, self.mode)
# set the file to the prev location.
fd.seek(self.loc)
self._fd = fd
def __getstate__(self):
with self._lock:
log.info("saving context state")
state = self.__dict__.copy()
for task in state['active'].values():
task.update(STATUS.INTERRUPTED)
state['pending'].append(task)
# close file descriptors on all active stages
for active_stage in state['active_stages'].values():
active_stage.close()
state['active'] = {}
# Remove the unpickable entries.
del state['_lock']
return state
def create(cls, sock):
try:
op_size = struct.calcsize(cls.HEADER)
op_bytes = sock.recv(op_size)
if len(op_bytes) == 0:
# agent disconnected !
return
op, = struct.unpack(cls.HEADER, op_bytes)
subs = cls.__subclasses__()
for operation in subs:
if operation.op_code.value == op:
return operation(sock=sock)
return None
except (struct.error, ValueError) as e:
log.info("Error parsing the message %s" % e)
return None
def create(cls, options):
"""
:param options: instance of `ServerConfig`
:type options: ´ServerConfig´
:return: instance of `Context`
:rtype: ´Context`
"""
rpath = options.resume_path
if os.path.isfile(rpath) and os.stat(rpath).st_size > 0:
log.info("Found resume file, loading...!")
with open(options.resume_path, 'rb') as rfile:
# i had to make this to make this testable with mocks!
# load with file didn't work, some how!
data = rfile.read()
ctx = pickle.loads(data)
return ctx
else:
return cls(options)
def handle(self):
"""
First method to be called by `BaseRequestHandler`.
responsible for initial call to authentication `do_auth`,
and `dispatcher`, the connection is kept alive as long as agent is
connected and their are targets to be delivered.
"""
log.info(f"{self.client_address} connected!")
self.connected = True
try:
while self.is_connected:
try:
# start by requesting authentication
if not self.authenticated:
self.do_auth()
self.dispatcher()
except (socket.timeout, ConnectionError) as e:
log.info(f"{self.client_address} Timeout - {e}")
self.connected = False
# mark any running task as interrupted
# so that other agent can take it later
self.ctx.interrupted(self.agent)
# wait a bit, in case a shutdown was requested!
self._terminate.wait(1.0)
finally:
if self.ctx.is_finished:
log.info("All stages are finished sending terminate event.")
self.server.shutdown()
self.request.close()
def _update_task_status(self, agent, status):
"""
Internal method updates a task of a given stage status, its also
responsible for managing the interrupted tasks.
:param agent: str with ipaddress and port in ip:port format
:param status: `STATUS` value to change.
"""
with self._lock:
task, tstage = self.__find_task_stage(agent)
if task and tstage:
task.update(status)
if status == STATUS.COMPLETED:
tstage.inc_finished()
# clean the completed task
del self.active[agent]
if status == status.INTERRUPTED:
log.info(f"Scan of {task.target} running on {agent} was "
f"interrupted")
self.pending.append(task)
del self.active[agent]
else:
log.debug(f"Agent {agent} is trying to update {status} on "
f"non existing task")
def do_report(self):
"""
When the scan the ends, the agent notifies the server that is ready
to send the report.
This method will handle the report transfer save the report in the
reports directory and make the target as finished
if the file hashes match.
"""
log.info("Agent Reporting Complete Scan!")
log.info(f"Filename {self.msg.filename} total file size "
f"{self.msg.filesize} file hash {self.msg.filehash}")
file_size = self.msg.filesize
nbytes = 0
report = self.ctx.get_report(self.agent,
self.msg.filename.decode("utf-8"))
try:
digest = hashlib.sha512()
self.ctx.downloading(self.agent)
while nbytes < file_size:
data = self.request.recv(1024)
report.write(data)
digest.update(data)
nbytes = nbytes + len(data)
if not hmac.compare_digest(digest.hexdigest().encode("utf-8"),
self.msg.filehash):
log.error(f"Files are not equal! {digest.hexdigest()}")
self.send_status(Status.FAILED)
else:
log.info("files are equal!")
self.ctx.completed(self.agent)
self.send_status(Status.SUCCESS)
finally:
if report:
report.flush()
report.close()
def __setstate__(self, state):
# restore the previous state, needed due to the existence of non
# serializable objects
self.__dict__.update(state)
log.info("Restoring context state")
self._lock = threading.Lock()
def do_ready(self):
"""
After the authentication the agent notifies the server, that is
ready to start scanning.
This will handle the request and send a target to be scanned.
"""
log.info("is Ready for targets")
log.info(f"Agent is running with uid {self.msg.uid}")
if self.msg.uid != 0:
log.info("Waning! agent is not running as root "
"syn scans might abort not enough privileges!")
target_data = self.ctx.pop(self.agent)
if not target_data:
if self.ctx.is_finished:
log.info("Target is None and all stages are finished")
# send empty command and terminate!
cmd = Command("", "")
self.request.sendall(cmd.pack())
self.connected = False
else:
log.info("Waiting for a stage to finish")
cmd = ExitStatus(Status.UNFINISHED)
self.request.sendall(cmd.pack())
return
cmd = Command(*target_data)
self.request.sendall(cmd.pack())
status_bytes = self.request.recv(1)
if len(status_bytes) == 0:
self.connected = False
log.info("Disconnected!")
self.ctx.interrupted(self.agent)
return
status, = struct.unpack("<B", status_bytes)
if status == Status.SUCCESS.value:
log.info("Started scanning !")
self.ctx.running(self.agent)
else:
log.error("Scan command returned Error")
log.info("Server is Terminating connection!")
self.connected = False
self.ctx.interrupted(self.agent)