Change port and protocol to annotated validators

This commit is contained in:
Marek Wolan
2024-09-25 16:28:22 +01:00
parent 695891f55c
commit f1b911bc65
84 changed files with 380 additions and 392 deletions

View File

@@ -7,9 +7,10 @@ from pydantic import validate_call
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.network.transmission.transport_layer import PORT_LOOKUP
from primaite.simulator.system.applications.application import Application
from primaite.utils.validators import IPV4Address, PROTOCOL_LOOKUP
from primaite.utils.validation.ip_protocol import IPProtocol, is_valid_protocol, PROTOCOL_LOOKUP
from primaite.utils.validation.ipv4_address import IPV4Address
from primaite.utils.validation.port import is_valid_port, Port, PORT_LOOKUP
class PortScanPayload(SimComponent):
@@ -23,8 +24,8 @@ class PortScanPayload(SimComponent):
"""
ip_address: IPV4Address
port: int
protocol: str
port: Port
protocol: IPProtocol
request: bool = True
def describe_state(self) -> Dict:
@@ -217,14 +218,14 @@ class NMAP(Application, identifier="NMAP"):
print(table.get_string(sortby="IP Address"))
return active_nodes
def _determine_port_scan_type(self, target_ip_addresses: List[IPV4Address], target_ports: List[int]) -> str:
def _determine_port_scan_type(self, target_ip_addresses: List[IPV4Address], target_ports: List[Port]) -> str:
"""
Determine the type of port scan based on the number of target IP addresses and ports.
:param target_ip_addresses: The list of target IP addresses.
:type target_ip_addresses: List[IPV4Address]
:param target_ports: The list of target ports.
:type target_ports: List[int]
:type target_ports: List[Port]
:return: The type of port scan.
:rtype: str
@@ -237,8 +238,8 @@ class NMAP(Application, identifier="NMAP"):
def _check_port_open_on_ip_address(
self,
ip_address: IPv4Address,
port: int,
protocol: str,
port: Port,
protocol: IPProtocol,
is_re_attempt: bool = False,
port_scan_uuid: Optional[str] = None,
) -> bool:
@@ -250,7 +251,7 @@ class NMAP(Application, identifier="NMAP"):
:param port: The target port.
:type port: Port
:param protocol: The protocol used for the port scan.
:type protocol: str
:type protocol: IPProtocol
:param is_re_attempt: Flag indicating if this is a reattempt. Defaults to False.
:type is_re_attempt: bool
:param port_scan_uuid: The UUID of the port scan payload. Defaults to None.
@@ -319,20 +320,20 @@ class NMAP(Application, identifier="NMAP"):
def port_scan(
self,
target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[int, List[int]]] = None,
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
target_port: Optional[Union[Port, List[Port]]] = None,
show: bool = True,
json_serializable: bool = False,
) -> Dict[IPv4Address, Dict[str, List[int]]]:
) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
"""
Perform a port scan on the target IP address(es).
:param target_ip_address: The target IP address(es) or network(s) for the port scan.
:type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
:param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
:type target_protocol: Optional[Union[str, List[str]]]
:type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
:param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
:type target_port: Optional[Union[int, List[int]]]
:type target_port: Optional[Union[Port, List[Port]]]
:param show: Flag indicating whether to display the scan results. Defaults to True.
:type show: bool
:param json_serializable: Flag indicating whether the return value should be JSON serializable. Defaults to
@@ -340,16 +341,16 @@ class NMAP(Application, identifier="NMAP"):
:type json_serializable: bool
:return: A dictionary mapping IP addresses to protocols and lists of open ports.
:rtype: Dict[IPv4Address, Dict[str, List[int]]]
:rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
"""
ip_addresses = self._explode_ip_address_network_array(target_ip_address)
if isinstance(target_port, int):
if is_valid_port(target_port):
target_port = [target_port]
elif target_port is None:
target_port = [port for port in PORT_LOOKUP if port not in {PORT_LOOKUP["NONE"], PORT_LOOKUP["UNUSED"]}]
if isinstance(target_protocol, str):
if is_valid_protocol(target_protocol):
target_protocol = [target_protocol]
elif target_protocol is None:
target_protocol = [PROTOCOL_LOOKUP["TCP"], PROTOCOL_LOOKUP["UDP"]]
@@ -389,12 +390,12 @@ class NMAP(Application, identifier="NMAP"):
def network_service_recon(
self,
target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]],
target_protocol: Optional[Union[str, List[str]]] = None,
target_port: Optional[Union[int, List[int]]] = None,
target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]] = None,
target_port: Optional[Union[Port, List[Port]]] = None,
show: bool = True,
show_online_only: bool = True,
json_serializable: bool = False,
) -> Dict[IPv4Address, Dict[str, List[int]]]:
) -> Dict[IPv4Address, Dict[IPProtocol, List[Port]]]:
"""
Perform a network service reconnaissance which includes a ping scan followed by a port scan.
@@ -407,9 +408,9 @@ class NMAP(Application, identifier="NMAP"):
:param target_ip_address: The target IP address(es) or network(s) for the port scan.
:type target_ip_address: Union[IPV4Address, List[IPV4Address], IPv4Network, List[IPv4Network]]
:param target_protocol: The protocol(s) to use for the port scan. Defaults to None, which includes TCP and UDP.
:type target_protocol: Optional[Union[str, List[str]]]
:type target_protocol: Optional[Union[IPProtocol, List[IPProtocol]]]
:param target_port: The port(s) to scan. Defaults to None, which includes all valid ports.
:type target_port: Optional[Union[int, List[int]]]
:type target_port: Optional[Union[Port, List[Port]]]
:param show: Flag indicating whether to display the scan results. Defaults to True.
:type show: bool
:param show_online_only: Flag indicating whether to show only the online hosts. Defaults to True.
@@ -419,7 +420,7 @@ class NMAP(Application, identifier="NMAP"):
:type json_serializable: bool
:return: A dictionary mapping IP addresses to protocols and lists of open ports.
:rtype: Dict[IPv4Address, Dict[str, List[int]]]
:rtype: Dict[IPv4Address, Dict[IPProtocol, List[Port]]]
"""
ping_scan_results = self.ping_scan(
target_ip_address=target_ip_address, show=show, show_online_only=show_online_only, json_serializable=False