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

#2248 - synced wth dev

Browse Source
This commit is contained in:
Chris McCarthy
2024-02-08 16:15:57 +00:00
parent 114fb8c077 c3ea06aeec
commit 5eed467271
20 changed files with 607 additions and 804 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -94,12 +94,12 @@ class FileSystem(SimComponent):
self._restore_manager.add_request(
name="file",
request_type=RequestType(
func=lambda request, context: self.restore_file(folder_uuid=request[0], file_uuid=request[1])
func=lambda request, context: self.restore_file(folder_name=request[0], file_name=request[1])
),
)
self._restore_manager.add_request(
name="folder",
request_type=RequestType(func=lambda request, context: self.restore_folder(folder_uuid=request[0])),
request_type=RequestType(func=lambda request, context: self.restore_folder(folder_name=request[0])),
)
rm.add_request(
name="restore",
@@ -209,7 +209,7 @@ class FileSystem(SimComponent):
folder = self.get_folder_by_id(folder_uuid=folder_uuid)
self.delete_folder(folder_name=folder.name)
def get_folder(self, folder_name: str) -> Optional[Folder]:
def get_folder(self, folder_name: str, include_deleted: bool = False) -> Optional[Folder]:
"""
Get a folder by its name if it exists.
@@ -219,9 +219,13 @@ class FileSystem(SimComponent):
for folder in self.folders.values():
if folder.name == folder_name:
return folder
if include_deleted:
for folder in self.deleted_folders.values():
if folder.name == folder_name:
return folder
return None
def get_folder_by_id(self, folder_uuid: str, include_deleted: bool = False) -> Optional[Folder]:
def get_folder_by_id(self, folder_uuid: str, include_deleted: Optional[bool] = False) -> Optional[Folder]:
"""
Get a folder by its uuid if it exists.
@@ -283,7 +287,7 @@ class FileSystem(SimComponent):
self._file_request_manager.add_request(name=file.name, request_type=RequestType(func=file._request_manager))
return file
def get_file(self, folder_name: str, file_name: str) -> Optional[File]:
def get_file(self, folder_name: str, file_name: str, include_deleted: Optional[bool] = False) -> Optional[File]:
"""
Retrieve a file by its name from a specific folder.
@@ -291,9 +295,9 @@ class FileSystem(SimComponent):
:param file_name: The name of the file to be retrieved, including its extension.
:return: An instance of File if it exists, otherwise `None`.
"""
folder = self.get_folder(folder_name)
folder = self.get_folder(folder_name, include_deleted=include_deleted)
if folder:
return folder.get_file(file_name)
return folder.get_file(file_name, include_deleted=include_deleted)
self.sys_log.info(f"File not found /{folder_name}/{file_name}")
def get_file_by_id(
@@ -455,46 +459,44 @@ class FileSystem(SimComponent):
for folder_id in self.folders:
self.folders[folder_id].reveal_to_red(instant_scan=instant_scan)
def restore_folder(self, folder_uuid: str):
def restore_folder(self, folder_name: str):
"""
Restore a folder.
Checks the current folder's status and applies the correct fix for the folder.
:param: folder_uuid: id of the folder to restore
:param: folder_name: name of the folder to restore
:type: folder_uuid: str
"""
folder = self.get_folder_by_id(folder_uuid=folder_uuid, include_deleted=True)
folder = self.get_folder(folder_name=folder_name, include_deleted=True)
if folder is None:
self.sys_log.error(f"Unable to restore folder with uuid {folder_uuid}. Folder does not exist.")
self.sys_log.error(f"Unable to restore folder {folder_name}. Folder is not in deleted folder list.")
return
self.deleted_folders.pop(folder.uuid, None)
folder.restore()
self.folders[folder.uuid] = folder
if folder.deleted:
self.deleted_folders.pop(folder.uuid)
def restore_file(self, folder_uuid: str, file_uuid: str):
def restore_file(self, folder_name: str, file_name: str):
"""
Restore a file.
Checks the current file's status and applies the correct fix for the file.
:param: folder_uuid: id of the folder where the file is stored
:type: folder_uuid: str
:param: folder_name: name of the folder where the file is stored
:type: folder_name: str
:param: file_uuid: id of the file to restore
:type: file_uuid: str
:param: file_name: name of the file to restore
:type: file_name: str
"""
folder = self.get_folder_by_id(folder_uuid=folder_uuid, include_deleted=True)
folder = self.get_folder(folder_name=folder_name)
if folder:
file = folder.get_file_by_id(file_uuid=file_uuid, include_deleted=True)
file = folder.get_file(file_name=file_name, include_deleted=True)
if file is None:
self.sys_log.error(f"Unable to restore file with uuid {file_uuid}. File does not exist.")
self.sys_log.error(f"Unable to restore file {file_name}. File does not exist.")
return
folder.restore_file(file_uuid=file_uuid)
folder.restore_file(file_name=file_name)

24
src/primaite/simulator/file_system/folder.py
View File

@@ -193,19 +193,19 @@ class Folder(FileSystemItemABC):
if self.restore_countdown == 0:
# repair all files
for file_id in self.files:
self.restore_file(file_uuid=file_id)
for file_id, file in self.files.items():
self.restore_file(file_name=file.name)
deleted_files = self.deleted_files.copy()
for file_id in deleted_files:
self.restore_file(file_uuid=file_id)
for file_id, file in deleted_files.items():
self.restore_file(file_name=file.name)
if self.deleted:
self.deleted = False
elif self.health_status in [FileSystemItemHealthStatus.CORRUPT, FileSystemItemHealthStatus.RESTORING]:
self.health_status = FileSystemItemHealthStatus.GOOD
def get_file(self, file_name: str) -> Optional[File]:
def get_file(self, file_name: str, include_deleted: Optional[bool] = False) -> Optional[File]:
"""
Get a file by its name.
@@ -218,6 +218,10 @@ class Folder(FileSystemItemABC):
for file in self.files.values():
if file.name == file_name:
return file
if include_deleted:
for file in self.deleted_files.values():
if file.name == file_name:
return file
return None
def get_file_by_id(self, file_uuid: str, include_deleted: Optional[bool] = False) -> File:
@@ -297,23 +301,23 @@ class Folder(FileSystemItemABC):
self.files = {}
def restore_file(self, file_uuid: str):
def restore_file(self, file_name: str):
"""
Restores a file.
:param file_uuid: The id of the file to restore
:param file_name: The name of the file to restore
"""
# if the file was not deleted, run a repair
file = self.get_file_by_id(file_uuid=file_uuid, include_deleted=True)
file = self.get_file(file_name=file_name, include_deleted=True)
if not file:
self.sys_log.error(f"Unable to restore file with uuid {file_uuid}. File does not exist.")
self.sys_log.error(f"Unable to restore file {file_name}. File does not exist.")
return
file.restore()
self.files[file.uuid] = file
if file.deleted:
self.deleted_files.pop(file_uuid)
self.deleted_files.pop(file.uuid)
def quarantine(self):
"""Quarantines the File System Folder."""

7
src/primaite/simulator/network/hardware/nodes/host/host_node.py
View File

@@ -153,13 +153,6 @@ class HostARP(ARP):
)
return
# Matched ARP request
# TODO: try taking this out
self.add_arp_cache_entry(
ip_address=arp_packet.sender_ip_address,
mac_address=arp_packet.sender_mac_addr,
network_interface=from_network_interface,
)
arp_packet = arp_packet.generate_reply(from_network_interface.mac_address)
self.send_arp_reply(arp_packet)

53
src/primaite/simulator/system/applications/web_browser.py
View File

@@ -1,7 +1,10 @@
from enum import Enum
from ipaddress import IPv4Address
from typing import Dict, Optional
from typing import Dict, List, Optional
from urllib.parse import urlparse
from pydantic import BaseModel, ConfigDict
from primaite import getLogger
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.network.protocols.http import (
@@ -33,6 +36,9 @@ class WebBrowser(Application):
latest_response: Optional[HttpResponsePacket] = None
"""Keeps track of the latest HTTP response."""
history: List["BrowserHistoryItem"] = []
"""Keep a log of visited websites and information about the visit, such as response code."""
def __init__(self, **kwargs):
kwargs["name"] = "WebBrowser"
kwargs["protocol"] = IPProtocol.TCP
@@ -71,7 +77,7 @@ class WebBrowser(Application):
:return: A dictionary capturing the current state of the WebBrowser and its child objects.
"""
state = super().describe_state()
state["last_response_status_code"] = self.latest_response.status_code if self.latest_response else None
state["history"] = [hist_item.state() for hist_item in self.history]
return state
def reset_component_for_episode(self, episode: int):
@@ -119,7 +125,8 @@ class WebBrowser(Application):
# create HTTPRequest payload
payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url=url)
# send request
# send request - As part of the self.send call, a response will be received and stored in the
# self.latest_response variable
if self.send(
payload=payload,
dest_ip_address=self.domain_name_ip_address,
@@ -129,9 +136,21 @@ class WebBrowser(Application):
f"{self.name}: Received HTTP {payload.request_method.name} "
f"Response {payload.request_url} - {self.latest_response.status_code.value}"
)
self.history.append(
WebBrowser.BrowserHistoryItem(
url=url,
status=self.BrowserHistoryItem._HistoryItemStatus.LOADED,
response_code=self.latest_response.status_code,
)
)
return self.latest_response.status_code is HttpStatusCode.OK
else:
self.sys_log.error(f"Error sending Http Packet {str(payload)}")
self.history.append(
WebBrowser.BrowserHistoryItem(
url=url, status=self.BrowserHistoryItem._HistoryItemStatus.SERVER_UNREACHABLE
)
)
return False
def send(
@@ -172,3 +191,31 @@ class WebBrowser(Application):
self.sys_log.info(f"{self.name}: Received HTTP {payload.status_code.value}")
self.latest_response = payload
return True
class BrowserHistoryItem(BaseModel):
"""Simple representation of browser history, used for tracking success of web requests to calculate rewards."""
model_config = ConfigDict(extra="forbid")
"""Error if incorrect specification."""
url: str
"""The URL that was attempted to be fetched by the browser"""
class _HistoryItemStatus(Enum):
NOT_SENT = "NOT_SENT"
PENDING = "PENDING"
SERVER_UNREACHABLE = "SERVER_UNREACHABLE"
LOADED = "LOADED"
status: _HistoryItemStatus = _HistoryItemStatus.PENDING
response_code: Optional[HttpStatusCode] = None
"""HTTP response code that was received, or PENDING if a response was not yet received."""
def state(self) -> Dict:
"""Return the contents of this dataclass as a dict for use with describe_state method."""
if self.status == self._HistoryItemStatus.LOADED:
outcome = self.response_code
else:
outcome = self.status.value
return {"url": self.url, "outcome": outcome}

4
src/primaite/simulator/system/services/database/database_service.py
View File

@@ -203,6 +203,10 @@ class DatabaseService(Service):
"""
self.sys_log.info(f"{self.name}: Running {query}")
if not self.db_file:
self.sys_log.info(f"{self.name}: Failed to run {query} because the database file is missing.")
return {"status_code": 404, "data": False}
if query == "SELECT":
if self.db_file.health_status == FileSystemItemHealthStatus.GOOD:
return {