oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#2682 Backport changes to core PrimAITE

Browse Source
This commit is contained in:
Marek Wolan
2024-06-25 11:04:52 +01:00
parent 4a81dc3b2c
commit 28dabad66b
52 changed files with 456 additions and 296 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/primaite/simulator/core.py
View File

@@ -221,7 +221,7 @@ class SimComponent(BaseModel):
}
return state
@validate_call
# @validate_call # this slows down execution quite a bit.
def apply_request(self, request: RequestFormat, context: Optional[Dict] = None) -> RequestResponse:
"""
Apply a request to a simulation component. Request data is passed in as a 'namespaced' list of strings.

17
src/primaite/simulator/file_system/file.py
View File

@@ -6,6 +6,8 @@ import json
import warnings
from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC, FileSystemItemHealthStatus
from primaite.simulator.file_system.file_type import FileType, get_file_type_from_extension
@@ -203,3 +205,18 @@ class File(FileSystemItemABC):
self.deleted = True
self.sys_log.info(f"File deleted {self.folder_name}/{self.name}")
return True
def show(self, markdown: bool = False):
"""
Prints a table of the file, displaying the file type, name and the file's access count.
:param markdown: Flag indicating if output should be in markdown format.
"""
headers = ["File Name", "File Type", "Number of Accesses"]
table = PrettyTable(headers)
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.name}"
table.add_row([self.name, self.file_type.name, self.num_access])
print(table)

101
src/primaite/simulator/file_system/file_system.py
View File

@@ -67,7 +67,7 @@ class FileSystem(SimComponent):
self._create_manager = RequestManager()
def _create_file_action(request: List[Any], context: Any) -> RequestResponse:
file = self.create_file(folder_name=request[0], file_name=request[1])
file = self.create_file(folder_name=request[0], file_name=request[1], force=request[2])
if not file:
return RequestResponse.from_bool(False)
return RequestResponse(
@@ -75,7 +75,7 @@ class FileSystem(SimComponent):
data={
"file_name": file.name,
"folder_name": file.folder_name,
"file_type": file.file_type,
"file_type": file.file_type.name,
"file_size": file.size,
},
)
@@ -111,9 +111,9 @@ class FileSystem(SimComponent):
data={
"file_name": file.name,
"folder_name": file.folder_name,
"file_type": file.file_type,
"file_type": file.file_type.name,
"file_size": file.size,
"file_status": file.health_status,
"file_status": file.health_status.name,
},
)
return RequestResponse.from_bool(False)
@@ -160,6 +160,21 @@ class FileSystem(SimComponent):
"""
return sum(folder.size for folder in self.folders.values())
def show_num_files(self, markdown: bool = False):
"""
Prints a table showing a host's number of file creations & deletions.
:param markdown: Flag indicating if output should be in markdown format.
"""
headers = ["File creations", "File deletions"]
table = PrettyTable(headers)
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
table.title = f"{self.sys_log.hostname} Number of Creations & Deletions"
table.add_row([self.num_file_creations, self.num_file_deletions])
print(table)
def show(self, markdown: bool = False, full: bool = False):
"""
Prints a table of the FileSystem, displaying either just folders or full files.
@@ -167,7 +182,7 @@ class FileSystem(SimComponent):
:param markdown: Flag indicating if output should be in markdown format.
:param full: Flag indicating if to show full files.
"""
headers = ["Folder", "Size", "Deleted"]
headers = ["Folder", "Size", "Health status", "Visible health status", "Deleted"]
if full:
headers[0] = "File Path"
table = PrettyTable(headers)
@@ -178,14 +193,38 @@ class FileSystem(SimComponent):
folders = {**self.folders, **self.deleted_folders}
for folder in folders.values():
if not full:
table.add_row([folder.name, folder.size_str, folder.deleted])
table.add_row(
[
folder.name,
folder.size_str,
folder.health_status.name,
folder.visible_health_status.name,
folder.deleted,
]
)
else:
files = {**folder.files, **folder.deleted_files}
if not files:
table.add_row([folder.name, folder.size_str, folder.deleted])
table.add_row(
[
folder.name,
folder.size_str,
folder.health_status.name,
folder.visible_health_status.name,
folder.deleted,
]
)
else:
for file in files.values():
table.add_row([file.path, file.size_str, file.deleted])
table.add_row(
[
file.path,
file.size_str,
file.health_status.name,
file.visible_health_status.name,
file.deleted,
]
)
if full:
print(table.get_string(sortby="File Path"))
else:
@@ -201,15 +240,15 @@ class FileSystem(SimComponent):
:param folder_name: The name of the folder.
"""
# check if folder with name already exists
if self.get_folder(folder_name):
raise Exception(f"Cannot create folder as it already exists: {folder_name}")
folder = Folder(name=folder_name, sys_log=self.sys_log)
folder = self.get_folder(folder_name)
if folder:
self.sys_log.info(f"Cannot create folder as it already exists: {folder_name}")
else:
folder = Folder(name=folder_name, sys_log=self.sys_log)
self._folder_request_manager.add_request(
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
self.folders[folder.uuid] = folder
self._folder_request_manager.add_request(
name=folder.name, request_type=RequestType(func=folder._request_manager)
)
return folder
def delete_folder(self, folder_name: str) -> bool:
@@ -289,6 +328,7 @@ class FileSystem(SimComponent):
size: Optional[int] = None,
file_type: Optional[FileType] = None,
folder_name: Optional[str] = None,
force: Optional[bool] = False,
) -> File:
"""
Creates a File and adds it to the list of files.
@@ -297,6 +337,7 @@ class FileSystem(SimComponent):
:param size: The size the file takes on disk in bytes.
:param file_type: The type of the file.
:param folder_name: The folder to add the file to.
:param force: Replaces the file if it already exists.
"""
if folder_name:
# check if file with name already exists
@@ -308,17 +349,23 @@ class FileSystem(SimComponent):
# Use root folder if folder_name not supplied
folder = self.get_folder("root")
# Create the file and add it to the folder
file = File(
name=file_name,
sim_size=size,
file_type=file_type,
folder_id=folder.uuid,
folder_name=folder.name,
sim_root=self.sim_root,
sys_log=self.sys_log,
)
folder.add_file(file)
file = self.get_file(folder, file_name)
if file:
self.sys_log.info(f"Cannot create file {file_name} as it already exists.")
if force:
self.sys_log.info(f"Replacing {file_name}")
else:
# Create the file and add it to the folder
file = File(
name=file_name,
sim_size=size,
file_type=file_type,
folder_id=folder.uuid,
folder_name=folder.name,
sim_root=self.sim_root,
sys_log=self.sys_log,
)
folder.add_file(file, force=force)
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
# increment file creation
self.num_file_creations += 1

21
src/primaite/simulator/file_system/file_type.py
View File

@@ -96,6 +96,20 @@ class FileType(Enum):
DB = 32
"Generic DB file. Used by sqlite3."
# Script file types
PS1 = 33
"A Powershell script. Used in Windows operating systems."
BAT = 34
"Windows batch file. Used in Windows operating systems."
SH = 35
"Linux shell script. Used in Unix based operating systems."
# Executable file types
PE = 36
"Portable Executable. Used in Windows operating systems."
ELF = 37
"Executable and Linkable Format. Used in Unix based operating systems."
@classmethod
def _missing_(cls, value: Any) -> FileType:
return cls.UNKNOWN
@@ -116,8 +130,11 @@ class FileType(Enum):
Returns 0 if a default size does not exist.
"""
size = file_type_sizes_bytes[self]
return size if size else 0
try:
size = file_type_sizes_bytes[self]
return size
except KeyError:
return 0
def get_file_type_from_extension(file_type_extension: str) -> FileType:

4
src/primaite/simulator/network/hardware/base.py
View File

@@ -11,6 +11,7 @@ from typing import Any, Dict, Optional, Type, TypeVar, Union
from prettytable import MARKDOWN, PrettyTable
from pydantic import BaseModel, Field
import primaite.simulator.network.nmne
from primaite import getLogger
from primaite.exceptions import NetworkError
from primaite.interface.request import RequestResponse
@@ -1049,7 +1050,7 @@ class Node(SimComponent):
def show_nic(self, markdown: bool = False):
"""Prints a table of the NICs on the Node."""
table = PrettyTable(["Port", "Type", "MAC Address", "Address", "Speed", "Status"])
table = PrettyTable(["Port", "Type", "MAC Address", "Address", "Speed", "Status", "NMNE"])
if markdown:
table.set_style(MARKDOWN)
table.align = "l"
@@ -1066,6 +1067,7 @@ class Node(SimComponent):
ip_address,
network_interface.speed,
"Enabled" if network_interface.enabled else "Disabled",
network_interface.nmne if primaite.simulator.network.nmne.CAPTURE_NMNE else "Disabled",
]
)
print(table)

5
src/primaite/simulator/network/networks.py
View File

@@ -185,7 +185,10 @@ def arcd_uc2_network() -> Network:
db_client_2.run()
web_browser_2 = client_2.software_manager.software.get("WebBrowser")
web_browser_2.target_url = "http://arcd.com/users/"
network.connect(endpoint_b=client_2.network_interface[1], endpoint_a=switch_2.network_interface[2])
network.connect(
endpoint_b=client_2.network_interface[1],
endpoint_a=switch_2.network_interface[2],
)
# Domain Controller
domain_controller = Server(

4
src/primaite/simulator/network/transmission/network_layer.py
View File

@@ -25,6 +25,10 @@ class IPProtocol(Enum):
ICMP = "icmp"
"""Internet Control Message Protocol."""
def model_dump(self) -> str:
"""Return as JSON-serialisable string."""
return self.name
class Precedence(Enum):
"""

4
src/primaite/simulator/network/transmission/transport_layer.py
View File

@@ -72,6 +72,10 @@ class Port(Enum):
POSTGRES_SERVER = 5432
"Postgres SQL Server."
def model_dump(self) -> str:
"""Return a json-serialisable string."""
return self.name
class UDPHeader(BaseModel):
"""

8
src/primaite/simulator/system/applications/nmap.py
View File

@@ -87,7 +87,9 @@ class NMAP(Application):
def _init_request_manager(self) -> RequestManager:
def _ping_scan_action(request: List[Any], context: Any) -> RequestResponse:
results = self.ping_scan(target_ip_address=request[0]["target_ip_address"], json_serializable=True)
results = self.ping_scan(
target_ip_address=request[0]["target_ip_address"], show=request[0]["show"], json_serializable=True
)
if not self._can_perform_network_action():
return RequestResponse.from_bool(False)
return RequestResponse(
@@ -311,7 +313,6 @@ class NMAP(Application):
self.sys_log.info(
f"{self.name}: Responding to port scan request for port {payload.port.value} "
f"({payload.port.name}) over {payload.protocol.name}",
True,
)
self.software_manager.send_payload_to_session_manager(payload=payload, session_id=session_id)
@@ -370,7 +371,8 @@ class NMAP(Application):
port_open = self._check_port_open_on_ip_address(ip_address=ip_address, port=port, protocol=protocol)
if port_open:
table.add_row([ip_address, port.value, port.name, protocol.name])
if show:
table.add_row([ip_address, port.value, port.name, protocol.name])
_ip_address = ip_address if not json_serializable else str(ip_address)
_protocol = protocol if not json_serializable else protocol.value
_port = port if not json_serializable else port.value

2
src/primaite/simulator/system/services/arp/arp.py
View File

@@ -143,7 +143,7 @@ class ARP(Service):
self.sys_log.info(f"Cannot send ARP request to a network address {str(target_ip_address)}")
return
if target_ip_address == outbound_network_interface.ip_network.broadcast_address:
self.sys_log.info(f"Cannot send ARP request to a broadcast address {str(target_ip_address)}")
self.sys_log.info(f"Cannot send ARP request to a broadcast addresss {str(target_ip_address)}")
return
self.sys_log.info(f"Sending ARP request from NIC {outbound_network_interface} for ip {target_ip_address}")

6
src/primaite/simulator/system/services/icmp/icmp.py
View File

@@ -61,7 +61,7 @@ class ICMP(Service):
if target_ip_address.is_loopback:
self.sys_log.info("Pinging loopback address")
return any(network_interface.enabled for network_interface in self.network_interfaces.values())
self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=True)
self.sys_log.info(f"Pinging {target_ip_address}:", to_terminal=False)
sequence, identifier = 0, None
while sequence < pings:
sequence, identifier = self._send_icmp_echo_request(target_ip_address, sequence, identifier, pings)
@@ -77,7 +77,7 @@ class ICMP(Service):
f"Received = {request_replies}, "
f"Lost = {pings - request_replies} ({(pings - request_replies) / pings * 100}% loss)"
)
self.sys_log.info(output, to_terminal=True)
self.sys_log.info(output, to_terminal=False)
return passed
@@ -167,7 +167,7 @@ class ICMP(Service):
f"bytes={len(frame.payload)}, "
f"time={time_str}, "
f"TTL={frame.ip.ttl}",
to_terminal=True,
to_terminal=False,
)
if not self.request_replies.get(frame.icmp.identifier):
self.request_replies[frame.icmp.identifier] = 0