from typing import List
import csle_collector.traffic_manager.traffic_manager_pb2_grpc
import csle_collector.traffic_manager.traffic_manager_pb2
import csle_collector.constants.constants as constants
[docs]def get_traffic_status(stub: csle_collector.traffic_manager.traffic_manager_pb2_grpc.TrafficManagerStub,
timeout=constants.GRPC.TIMEOUT_SECONDS) \
-> csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO:
"""
Queries the server for the Traffic generator status
:param stub: the stub to send the remote gRPC to the server
:param timeout: the GRPC timeout (seconds)
:return: a TrafficDTO describing the status of the traffic generator
"""
get_traffic_status_msg = csle_collector.traffic_manager.traffic_manager_pb2.GetTrafficStatusMsg()
traffic_dto: csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO = \
stub.getTrafficStatus(get_traffic_status_msg, timeout=timeout)
return traffic_dto
[docs]def stop_traffic(stub: csle_collector.traffic_manager.traffic_manager_pb2_grpc.TrafficManagerStub,
timeout=constants.GRPC.TIMEOUT_SECONDS) \
-> csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO:
"""
Sends a request to the traffic manager to stop the traffic generator
:param stub: the stub to send the remote gRPC to the server
:param timeout: the GRPC timeout (seconds)
:return: a TrafficDTO describing the status of the traffic generator
"""
stop_traffic_msg = csle_collector.traffic_manager.traffic_manager_pb2.StopTrafficMsg()
traffic_dto: csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO = \
stub.stopTraffic(stop_traffic_msg, timeout=timeout)
return traffic_dto
[docs]def start_traffic(stub: csle_collector.traffic_manager.traffic_manager_pb2_grpc.TrafficManagerStub,
commands: List[str], sleep_time: int,
timeout=constants.GRPC.TIMEOUT_SECONDS) \
-> csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO:
"""
Sends a request to the traffic manager to start the traffic generator script
:param stub: the stub to send the remote gRPC to the server
:param timeout: the GRPC timeout (seconds)
:param commands: list of commands for the traffic generator
:param sleep_time: sleep time for the traffic generator
:return: an TrafficDTO describing the status of the traffic manager
"""
start_traffic_msg = csle_collector.traffic_manager.traffic_manager_pb2.StartTrafficMsg(
commands=commands, sleepTime=sleep_time)
traffic_dto: csle_collector.traffic_manager.traffic_manager_pb2.TrafficDTO = \
stub.startTraffic(start_traffic_msg, timeout=timeout)
return traffic_dto