seb/2024-bsc-sebastian-lenzlinger
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Factor out pydantic.

Browse Source
This commit is contained in:
Sebastian Lenzlinger
2024-05-07 22:48:53 +02:00
parent 95426e0baa
commit 11e2c356fa
6 changed files with 55 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
code/iottb/logger.py
View File

@@ -5,16 +5,16 @@ from logging.handlers import RotatingFileHandler
def setup_logging():
logger_obj = logging.getLogger('iottbLogger')
logger_obj.setLevel(logging.INFO)
logger_obj.setLevel(logging.DEBUG)
file_handler = RotatingFileHandler('iottb.log')
console_handler = logging.StreamHandler(sys.stdout)
file_handler.setLevel(logging.DEBUG)
console_handler.setLevel(logging.INFO)
file_handler.setLevel(logging.INFO)
console_handler.setLevel(logging.DEBUG)
file_fmt = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
console_fmt = logging.Formatter('%(name)s - %(levelname)s - %(message)s')
file_fmt = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
console_fmt = logging.Formatter('%(name)s - %(name)s - %(levelname)s - %(message)s')
file_handler.setFormatter(file_fmt)
console_handler.setFormatter(console_fmt)

22
code/iottb/models/capture_metadata_model.py
View File

@@ -5,19 +5,19 @@ from pathlib import Path
from typing import Optional, Any
from uuid import UUID
from pydantic import BaseModel, Field
from iottb.definitions import ReturnCodes, CAPTURE_METADATA_FILE
from iottb.models.device_metadata_model import DeviceMetadata
from iottb.logger import logger
class CaptureMetadata(BaseModel):
class CaptureMetadata:
# Required Fields
device_metadata: DeviceMetadata = Field(exclude=True)
capture_id: uuid.UUID = Field(default_factory=lambda: str(uuid.uuid4()))
device_metadata: DeviceMetadata
capture_id: uuid.UUID = lambda: str(uuid.uuid4())
capture_dir: Path
capture_file: str
capture_date: str = Field(default_factory=lambda: datetime.now().strftime('%d-%m-%YT%H:%M:%S').lower())
capture_date: str = lambda: datetime.now().strftime('%d-%m-%YT%H:%M:%S').lower()
# Statistics
start_time: str
@@ -38,14 +38,15 @@ class CaptureMetadata(BaseModel):
firmware_version: Optional[str] = None
def __init__(self, device_metadata: DeviceMetadata, capture_dir: Path, /, **data: Any):
logger.info(f"Creating CaptureMetadata model from DeviceMetadata: {device_metadata}")
super().__init__(**data) # Pycharms orders
self.device_metadata = device_metadata
self.capture_dir = capture_dir
assert capture_dir.is_dir()
assert capture_dir.is_dir(), f"Capture directory {capture_dir} does not exist"
# Getters
def get_device_id(self) -> str:
return self.device_id
return self.device_metadata.get_device_id()
def get_start_time(self) -> str:
return self.start_time
@@ -136,14 +137,18 @@ class CaptureMetadata(BaseModel):
# Other
def build_capture_file_name(self):
logger.info(f"Building capture file name")
prefix = ""
if self.app is None:
logger.debug(f"No app specified")
prefix = self.device_metadata.get_device_short_name()
else:
logger.debug(f"App specified: {self.app}")
assert str(self.app).strip() not in {"", " "}, f"app is not a valid name: {self.app}"
prefix = self.get_app()
# assert self.capture_dir is not None, f"{self.capture_dir} does not exist"
filename = f"{prefix}_{str(self.capture_id)}.pcap"
logger.debug(f"Capture file name: {filename}")
self.set_capture_file(filename)
def save_capture_metadata_to_json(self, file_path: Path = Path(CAPTURE_METADATA_FILE)):
@@ -155,3 +160,6 @@ class CaptureMetadata(BaseModel):
with file_path.open('w') as file:
json.dump(metadata, file)
return ReturnCodes.SUCCESS
def model_dump_json(self, indent, exclude_unset, exclude_none):
pass

37
code/iottb/models/device_metadata_model.py
View File

@@ -6,22 +6,22 @@ from typing import Optional, List, Any
# iottb modules
from iottb.definitions import ReturnCodes, DEVICE_METADATA_FILE
from iottb.logger import logger
# 3rd party libs
from pydantic import BaseModel, Field
IMMUTABLE_FIELDS = {"device_name", "device_short_name", "device_id", "date_created"}
class DeviceMetadata(BaseModel):
class DeviceMetadata:
# Required fields
device_name: str
device_short_name: str
device_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
date_created: str = Field(default_factory=lambda: datetime.now().strftime('%d-%m-%YT%H:%M:%S').lower())
device_id: str = lambda: str(uuid.uuid4())
date_created: str = lambda: datetime.now().strftime('%d-%m-%YT%H:%M:%S').lower()
device_root_path: Path
# Optional Fields
aliases: List[str] = Field(default_factory=lambda: [])
aliases: Optional[List[str]] = None
device_type: Optional[str] = None
device_serial_number: Optional[str] = None
device_firmware_version: Optional[str] = None
@@ -29,13 +29,16 @@ class DeviceMetadata(BaseModel):
capture_files: Optional[List[str]] = []
def __init__(self, device_name: str, device_root_dir: Path, /, **data: Any):
super().__init__(**data)
def __init__(self, device_name: str, device_root_dir: Path):
self.device_name = device_name
self.device_short_name = device_name.lower().replace(" ", "_")
# assert dir_contains_device_metadata(device_root_dir), \
# f"Directory {device_root_dir} is missing a {DEVICE_METADATA_FILE} file"
self.device_root_dir = device_root_dir
logger.debug(f"Device name: {device_name}")
logger.debug(f"Device short_name: {self.device_short_name}")
logger.debug(f"Device root dir: {device_root_dir}")
logger.info(f"Initialized DeviceMetadata model: {device_name}")
def get_device_id(self) -> str:
return self.device_id
@@ -83,15 +86,18 @@ class DeviceMetadata(BaseModel):
@classmethod
def load_from_json(cls, device_file_path: Path):
logger.info(f"Loading DeviceMetadata from JSON file: {device_file_path}")
assert device_file_path.is_file(), f"{device_file_path} is not a file"
assert device_file_path.name == DEVICE_METADATA_FILE, f"{device_file_path} is not a {DEVICE_METADATA_FILE}"
device_meta_filename = device_file_path
with device_meta_filename.open('r') as file:
metadata_json = json.load(file)
metadata_model_obj = cls.model_validate_json(metadata_json)
return metadata_model_obj
def save_to_json(self, file_path: Path):
logger.info(f"Saving DeviceMetadata to JSON file: {file_path}")
if file_path.is_file():
print(f"File {file_path} already exists, update instead.")
return ReturnCodes.FILE_ALREADY_EXISTS
@@ -100,19 +106,12 @@ class DeviceMetadata(BaseModel):
json.dump(metadata, file)
return ReturnCodes.SUCCESS
@classmethod
def update_metadata_in_json(cls, file_path: Path, **kwargs):
# TODO Maybe not needed at all.
assert file_path.is_file()
for field in IMMUTABLE_FIELDS:
if field in kwargs:
print(f"Field {field} is immutable")
return ReturnCodes.IMMUTABLE
metadata = cls.load_from_json(file_path)
for field, value in kwargs.items():
if field in metadata.model_fields_set:
setattr(metadata, field, value)
metadata.date_updated = datetime.now().strftime('%d-%m-%YT%H:%M:%S').lower()
def model_validate_json(cls, metadata_json):
pass
def model_dump_json(self, indent):
pass

17
code/iottb/subcommands/add_device.py
View File

@@ -2,6 +2,7 @@ import pathlib
from iottb import definitions
from iottb.definitions import DEVICE_METADATA_FILE, ReturnCodes
from iottb.logger import logger
from iottb.models.device_metadata_model import DeviceMetadata
from iottb.utils.device_metadata_utils import *
@@ -16,14 +17,15 @@ def setup_init_device_root_parser(subparsers):
def handle_add(args):
print("Entered add-device-root")
logger.info(f"Add device handler called with args {args}")
if args.guided:
logger.debug("Guided setup")
metadata = guided_setup(args.root_dir)
else:
logger.debug("Setup through passed args: setup")
device_name = args.name
args.root_dir.mkdir(parents=True, exist_ok=True)
args.root_dir.chdir()
metadata = DeviceMetadata(device_name, args.root_dir)
file_path = args.root_dir / DEVICE_METADATA_FILE
@@ -31,7 +33,7 @@ def handle_add(args):
if response.lower() not in definitions.AFFIRMATIVE_USER_RESPONSE.add(""):
configure_metadata()
assert False, "TODO implement dynamic setup"
assert metadata.model_dump() != ""
if metadata.save_to_json(file_path) == ReturnCodes.FILE_ALREADY_EXISTS:
print("Directory already contains a device metadata file. Aborting operation.")
return ReturnCodes.ABORTED
@@ -44,16 +46,17 @@ def configure_metadata():
def guided_setup(device_root) -> DeviceMetadata:
logger.info("Guided setup")
response = "N"
device_name = ""
while response.upper() == "N":
device_name = input("Please enter name of device: ")
response = input(f"Confirm device name: {device_name} [y/N] ")
if device_name == "" or device_name is None:
print("Name cannot be empty")
response = input(f"Confirm device name: {device_name} [y/N] ")
logger.warning("Name cannot be empty")
logger.debug(f"Response is {response}")
logger.debug(f"Device name is {device_name}")
assert response.lower() in definitions.AFFIRMATIVE_USER_RESPONSE.add(""), f"{response.upper()} not supported"
assert device_name != ""
assert device_name is not None
return DeviceMetadata(device_name, device_root)