Make all requests return a RequestResponse

This commit is contained in:
Marek Wolan
2024-03-08 14:58:34 +00:00
parent d331224b45
commit beb51834f9
13 changed files with 275 additions and 136 deletions

View File

@@ -7,6 +7,7 @@ from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.file_system.file import File
from primaite.simulator.file_system.file_type import FileType
@@ -41,12 +42,16 @@ class FileSystem(SimComponent):
self._delete_manager.add_request(
name="file",
request_type=RequestType(
func=lambda request, context: self.delete_file(folder_name=request[0], file_name=request[1])
func=lambda request, context: RequestResponse.from_bool(
self.delete_file(folder_name=request[0], file_name=request[1])
)
),
)
self._delete_manager.add_request(
name="folder",
request_type=RequestType(func=lambda request, context: self.delete_folder(folder_name=request[0])),
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(self.delete_folder(folder_name=request[0]))
),
)
rm.add_request(
name="delete",
@@ -57,12 +62,16 @@ class FileSystem(SimComponent):
self._restore_manager.add_request(
name="file",
request_type=RequestType(
func=lambda request, context: self.restore_file(folder_name=request[0], file_name=request[1])
func=lambda request, context: RequestResponse(
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_name=request[0])),
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(self.restore_folder(folder_name=request[0]))
),
)
rm.add_request(
name="restore",
@@ -138,7 +147,7 @@ class FileSystem(SimComponent):
)
return folder
def delete_folder(self, folder_name: str):
def delete_folder(self, folder_name: str) -> bool:
"""
Deletes a folder, removes it from the folders list and removes any child folders and files.
@@ -146,24 +155,26 @@ class FileSystem(SimComponent):
"""
if folder_name == "root":
self.sys_log.warning("Cannot delete the root folder.")
return
return False
folder = self.get_folder(folder_name)
if folder:
# set folder to deleted state
folder.delete()
# remove from folder list
self.folders.pop(folder.uuid)
# add to deleted list
folder.remove_all_files()
self.deleted_folders[folder.uuid] = folder
self.sys_log.info(f"Deleted folder /{folder.name} and its contents")
else:
if not folder:
_LOGGER.debug(f"Cannot delete folder as it does not exist: {folder_name}")
return False
def delete_folder_by_id(self, folder_uuid: str):
# set folder to deleted state
folder.delete()
# remove from folder list
self.folders.pop(folder.uuid)
# add to deleted list
folder.remove_all_files()
self.deleted_folders[folder.uuid] = folder
self.sys_log.info(f"Deleted folder /{folder.name} and its contents")
return True
def delete_folder_by_id(self, folder_uuid: str) -> None:
"""
Deletes a folder via its uuid.
@@ -297,7 +308,7 @@ class FileSystem(SimComponent):
return file
def delete_file(self, folder_name: str, file_name: str):
def delete_file(self, folder_name: str, file_name: str) -> bool:
"""
Delete a file by its name from a specific folder.
@@ -309,8 +320,10 @@ class FileSystem(SimComponent):
file = folder.get_file(file_name)
if file:
folder.remove_file(file)
return True
return False
def delete_file_by_id(self, folder_uuid: str, file_uuid: str):
def delete_file_by_id(self, folder_uuid: str, file_uuid: str) -> None:
"""
Deletes a file via its uuid.
@@ -327,7 +340,7 @@ class FileSystem(SimComponent):
else:
self.sys_log.error(f"Unable to delete file that does not exist. (id: {file_uuid})")
def move_file(self, src_folder_name: str, src_file_name: str, dst_folder_name: str):
def move_file(self, src_folder_name: str, src_file_name: str, dst_folder_name: str) -> None:
"""
Move a file from one folder to another.
@@ -404,7 +417,7 @@ class FileSystem(SimComponent):
# Agent actions
###############################################################
def scan(self, instant_scan: bool = False):
def scan(self, instant_scan: bool = False) -> None:
"""
Scan all the folders (and child files) in the file system.
@@ -413,7 +426,7 @@ class FileSystem(SimComponent):
for folder_id in self.folders:
self.folders[folder_id].scan(instant_scan=instant_scan)
def reveal_to_red(self, instant_scan: bool = False):
def reveal_to_red(self, instant_scan: bool = False) -> None:
"""
Reveals all the folders (and child files) in the file system to the red agent.
@@ -422,7 +435,7 @@ 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_name: str):
def restore_folder(self, folder_name: str) -> bool:
"""
Restore a folder.
@@ -435,13 +448,14 @@ class FileSystem(SimComponent):
if folder is None:
self.sys_log.error(f"Unable to restore folder {folder_name}. Folder is not in deleted folder list.")
return
return False
self.deleted_folders.pop(folder.uuid, None)
folder.restore()
self.folders[folder.uuid] = folder
return True
def restore_file(self, folder_name: str, file_name: str):
def restore_file(self, folder_name: str, file_name: str) -> bool:
"""
Restore a file.
@@ -454,12 +468,15 @@ class FileSystem(SimComponent):
:type: file_name: str
"""
folder = self.get_folder(folder_name=folder_name)
if not folder:
_LOGGER.debug(f"Cannot restore file {file_name} in folder {folder_name} as the folder does not exist.")
return False
if folder:
file = folder.get_file(file_name=file_name, 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 {file_name}. File does not exist.")
return
if not file:
msg = f"Unable to restore file {file_name}. File was not found."
self.sys_log.error(msg)
return False
folder.restore_file(file_name=file_name)
return folder.restore_file(file_name=file_name)

View File

@@ -6,6 +6,7 @@ from enum import Enum
from typing import Dict, Optional
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType, SimComponent
from primaite.simulator.system.core.sys_log import SysLog
@@ -102,12 +103,26 @@ class FileSystemItemABC(SimComponent):
def _init_request_manager(self) -> RequestManager:
rm = super()._init_request_manager()
rm.add_request(name="scan", request_type=RequestType(func=lambda request, context: self.scan()))
rm.add_request(name="checkhash", request_type=RequestType(func=lambda request, context: self.check_hash()))
rm.add_request(name="repair", request_type=RequestType(func=lambda request, context: self.repair()))
rm.add_request(name="restore", request_type=RequestType(func=lambda request, context: self.restore()))
rm.add_request(
name="scan", request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.scan()))
)
rm.add_request(
name="checkhash",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.check_hash())),
)
rm.add_request(
name="repair",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.repair())),
)
rm.add_request(
name="restore",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.restore())),
)
rm.add_request(name="corrupt", request_type=RequestType(func=lambda request, context: self.corrupt()))
rm.add_request(
name="corrupt",
request_type=RequestType(func=lambda request, context: RequestResponse.from_bool(self.corrupt())),
)
return rm
@@ -124,9 +139,9 @@ class FileSystemItemABC(SimComponent):
return convert_size(self.size)
@abstractmethod
def scan(self) -> None:
def scan(self) -> bool:
"""Scan the folder/file - updates the visible_health_status."""
pass
return False
@abstractmethod
def reveal_to_red(self) -> None:
@@ -134,7 +149,7 @@ class FileSystemItemABC(SimComponent):
pass
@abstractmethod
def check_hash(self) -> None:
def check_hash(self) -> bool:
"""
Checks the has of the file to detect any changes.
@@ -142,30 +157,30 @@ class FileSystemItemABC(SimComponent):
Return False if corruption is detected, otherwise True
"""
pass
return False
@abstractmethod
def repair(self) -> None:
def repair(self) -> bool:
"""
Repair the FileSystemItem.
True if successfully repaired. False otherwise.
"""
pass
return False
@abstractmethod
def corrupt(self) -> None:
def corrupt(self) -> bool:
"""
Corrupt the FileSystemItem.
True if successfully corrupted. False otherwise.
"""
pass
return False
@abstractmethod
def restore(self) -> None:
def restore(self) -> bool:
"""Restore the file/folder to the state before it got ruined."""
pass
return False
@abstractmethod
def delete(self) -> None:

View File

@@ -5,6 +5,7 @@ from typing import Dict, Optional
from prettytable import MARKDOWN, PrettyTable
from primaite import getLogger
from primaite.interface.request import RequestResponse
from primaite.simulator.core import RequestManager, RequestType
from primaite.simulator.file_system.file import File
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC, FileSystemItemHealthStatus
@@ -53,7 +54,9 @@ class Folder(FileSystemItemABC):
rm = super()._init_request_manager()
rm.add_request(
name="delete",
request_type=RequestType(func=lambda request, context: self.remove_file_by_id(file_uuid=request[0])),
request_type=RequestType(
func=lambda request, context: RequestResponse.from_bool(self.remove_file_by_name(file_name=request[0]))
),
)
self._file_request_manager = RequestManager()
rm.add_request(
@@ -249,6 +252,21 @@ class Folder(FileSystemItemABC):
file = self.get_file_by_id(file_uuid=file_uuid)
self.remove_file(file=file)
def remove_file_by_name(self, file_name: str) -> bool:
"""
Remove a file using its name.
:param file_name: filename
:type file_name: str
:return: Whether it was successfully removed.
:rtype: bool
"""
for f in self.files.values():
if f.name == file_name:
self.remove_file(f)
return True
return False
def remove_all_files(self):
"""Removes all the files in the folder."""
for file_id in self.files:
@@ -258,7 +276,7 @@ class Folder(FileSystemItemABC):
self.files = {}
def restore_file(self, file_name: str):
def restore_file(self, file_name: str) -> bool:
"""
Restores a file.
@@ -268,13 +286,14 @@ class Folder(FileSystemItemABC):
file = self.get_file(file_name=file_name, include_deleted=True)
if not file:
self.sys_log.error(f"Unable to restore file {file_name}. File does not exist.")
return
return False
file.restore()
self.files[file.uuid] = file
if file.deleted:
self.deleted_files.pop(file.uuid)
return True
def quarantine(self):
"""Quarantines the File System Folder."""