Source code for dscan.client

#!/usr/bin/env python3
# encoding: utf-8

"""
client.py
client side responsible for the managing clients and scan execution flow.
"""
import os
import hmac
import struct
import random
import threading
import time
from socket import socket
from socket import AF_INET
from socket import SOCK_STREAM
from socket import timeout
import ssl
from dscan import log
from dscan.models.structures import Structure, Operations
from dscan.models.structures import Auth
from dscan.models.structures import Ready
from dscan.models.structures import Status
from dscan.models.scanner import ScanProcess
from string import ascii_uppercase


[docs]class Agent: """ Agent client implementation. """ def __init__(self, config): """ :param config: `dscan.models.scanner.Config` instance with the runtime configurations. """ self.connected = False self.config = config self.con_retries = 0 srv_hostname = config.srv_hostname ssl_context = ssl.create_default_context() ssl_context.check_hostname = False ssl_context.load_verify_locations(self.config.sslcert) self.socket = ssl_context.wrap_socket(socket(AF_INET, SOCK_STREAM), server_side=False, server_hostname=srv_hostname) self._terminate = threading.Event() self.scan = ScanProcess(self.config.outdir)
[docs] def is_connected(self): """ Check if the agent is still connected and not yet finished. :return: `True` if the client has disconnected or the terminate event has been triggered, else False. :rtype: `bool` """ return self.con_retries < 3 and not self._terminate.is_set()
[docs] def start(self): """ Start the client connects to the server and authenticates. :return: True if was able to connect and authentication was successful else False. :rtype: `bool` """ self.con_retries = 0 # while the connection retry is under 3 tries # everytime the connection is interrupted the client # tries to connect authenticates and requests a target! while self.is_connected(): try: self.socket.connect((self.config.host, self.config.port)) self.connected = True if not self.do_auth(): # return out if auth was not successfully self.connected = False return # reset the counter if connection was successful. self.con_retries = 0 # if authentication was successful request a target to scan. self.do_ready() except (timeout, ConnectionError, ValueError) as e: self.con_retries += 1 log.error(f"Connection Timeout - {e}") log.error(f"Attempt - {self.con_retries} " f"to establish connection") self.connected = False self.socket.close() finally: self.socket.close()
[docs] def shutdown(self): """ Set the terminate event On, to shutdown the agent. """ self._terminate.set()
[docs] def do_auth(self): """ Initiate the authentication. :return: True if the status code of the last operation is `dscan.models.structures.Status.SUCCESS` False otherwise. :rtype: `bool` """ log.info("Initiating authentication") opr = Structure.create(self.socket) if not opr: # unable to get message return False hmac_hash = hmac.new(self.config.secret_key, opr.data, 'sha512') digest = hmac_hash.hexdigest().encode("utf-8") self.socket.sendall(Auth(digest).pack()) status_result = self.__check_status() if status_result: return status_result
[docs] def do_ready(self): """ This is recursive method and is responsible for, notifying the server is ready to execute a new scan, launch the scan and save the report. Until the server returns no target to scan. """ alias = "".join(random.choice(ascii_uppercase) for _ in range(6)) while self.connected: log.info("Requesting target...") self.socket.sendall(Ready(os.getuid(), alias).pack()) cmd = Structure.create(self.socket) if not cmd: # unable to get message log.info("Unable to receive command from server") return if cmd.op_code == Operations.STATUS and cmd.status == Status.FINISHED: log.info("received a Finished status, Terminating!") self.con_retries = 3 return if cmd.op_code == Operations.STATUS and cmd.status == Status.UNFINISHED: log.info("received a Unfinished will retry later!") time.sleep(5) log.info("retrying.. Target request!") continue log.info(f"Launching scan on {cmd}") report = self.scan.run(cmd.target.decode("utf-8"), cmd.options.decode("utf-8"), self.send_status) if report: self.socket.sendall(report.pack()) if self.__send_report(report): log.info("Report Transfer was successful") else: log.error("Report Transfer was unsuccessful") else: self.send_status(Status.FAILED.value)
[docs] def send_status(self, status): """ :param status: int of a valid `dscan.models.structures.Status` """ self.socket.sendall(struct.pack("<B", status))
def __check_status(self): """ Receives the status code from the server, and check the code value, see `dscan.models.structures.Status` for other valid status values. :return: True if the status code of the last operation is `dscan.models.structures.Status.SUCCESS` False otherwise. :rtype: `bool` """ op_size = struct.calcsize("<B") op_bytes = self.socket.recv(op_size) if len(op_bytes) == 0: log.info("disconnected!") self.connected = False return False status, = struct.unpack("<B", op_bytes) if status == Status.SUCCESS.value: log.info("Operation Successful ...") return True else: log.error("Operation unsuccessful disconnecting...") return False def __send_report(self, report, retry=0): """ Transfer the report. :param report: report message. :type report: `dscan.models.structures.Report` :param retry: number of attempts :type retry: `int` :return: bool the result of last try. :rtype: `bool` """ nbytes = 0 with open(os.path.join(self.config.outdir, report.filename.decode("utf-8")), "rb") as rfile: while nbytes < report.filesize: data = rfile.read(1024) self.socket.sendall(data) nbytes = nbytes + len(data) result = self.__check_status() if not result and retry < 3: retry += 1 log.info(f"Transfer retry {retry}") return self.__send_report(report, retry) else: return result