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

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