locust script to run scan command with the staging atsign

locust_script.py

from locust import User, task, between, events
import socket
import time
import statistics
import os
import sys
import logging
import ssl

#TcpClient Class:
#This class represents a TCP client that connects to a specified host and port.
#It initializes a socket and provides methods for connecting, closing the connection, and sending data.
#The send_data method sends data to the server and receives a response.

# Configure logging
logging.basicConfig(filename='my_log_file.log', level=logging.DEBUG)
logger = logging.getLogger()
logger.addHandler(logging.StreamHandler(sys.stdout))

class TcpClient:
    def __init__(self, host, port):
        self.host = host
        self.port = port
        self.context = ssl.create_default_context()
        self.sock = self.context.wrap_socket(socket.socket(socket.AF_INET, socket.SOCK_STREAM), server_hostname=host)
        self.connect()

    def connect(self):
        self.sock.connect((self.host, self.port))

    def close(self):
        self.sock.close()

    def send_data(self, data):
        try:
            self.sock.send(data.encode())
            response = self.sock.recv(1024).decode()
            return response
        except (socket.error, ssl.SSLError) as e:
            print(f"Error during send_data: {e}")
            # Handle the exception, such as by reconnecting or logging the error.
            # For simplicity, I'm closing and reconnecting the socket here.
            self.close()
            self.connect()
            return "Error: Connection reset"

#MyTaskSet Class (Locust User):
#Inherits from User in Locust.
#Defines the behavior of a user during load testing.
#The on_start method initializes the TCP client and authenticates.
#The custom_command task simulates sending a custom command to the server.
#Response time is tracked using Locust events.
#The on_stop method closes the client connection.

class MyTaskSet(User):
    wait_time = between(1, 3)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client = None  # Initialize client attribute to None

    def on_start(self):

        try:
            atSign = os.environ.get("ATSIGN", "default_atsign")
            print("atSign:", atSign)
            host = os.environ.get("HOST", "1bb5c229-db92-5e28-aa6b-b69ddd6553e4.staging0001.atsign.zone")
            print("host:", host)
            port = os.environ.get("PORT", "2035")
            print("port:", port)
            self.client = TcpClient(host, int(port))
            self.authenticate()
        except Exception as e:
            print(f"An error occurred during on_start: {e}")



    def on_stop(self):
        self.client.close()

    def authenticate(self):
        auth_message = "AUTHENTICATE\n"
        #self.client.send_data(auth_message)


    @task
    def custom_command(self):
        try:
            start_time = time.time()

            command_message = "scan\n"
            response = self.client.send_data(command_message)
            print("response:", response)

            end_time = time.time()
            response_time = int((end_time - start_time) * 1000)  # Convert to milliseconds

            # Track response time using Locust events
            events.request.fire(
                request_type="verb_execution",
                name="verb_execution",
                response_time=response_time,
                response_length=len(response),
            )

            print(f"Custom command response: {response} {response_time}")
        except Exception as e:
            print(f"An error occurred during on_start: {e}")

#Statistics and Event Handling:
#The on_locust_exit function prints statistics at the end of the test.
#It displays requests per second, total requests, failures, and the number of users.
def on_locust_exit(environment, **kwargs):
    stats = environment.runner.stats.total
    print("Test finished. Collecting and printing statistics...")

    # General statistics
    print(f"Requests per second: {stats.total_rps}")
    print(f"Total Requests: {stats.total_num_requests}")
    print(f"Total Failures: {stats.total_num_failures}")
    print(f"Total Users: {stats.total_user_count}")

    # Response time statistics
    print(f"Minimum Response Time: {stats.total_min_response_time} ms")
    print(f"Maximum Response Time: {stats.total_max_response_time} ms")
    print(f"Average Response Time: {stats.total_avg_response_time} ms")

    # Percentile calculations
    response_times = stats.response_times.values()
    p90 = statistics.percentile(response_times, 90)
    p99 = statistics.percentile(response_times, 99)
    print(f"90th Percentile Response Time: {p90} ms")
    print(f"99th Percentile Response Time: {p99} ms")

if __name__ == "__main__":
    events.locust_exit += on_locust_exit
    MyTaskSet().run()