Skip to content

Commit

Permalink
feat: Add support for META file
Browse files Browse the repository at this point in the history 
For this use a new entry point for commands and unify the exception
handling.
  • Loading branch information
@rumpelsepp
rumpelsepp committed Aug 25, 2022
1 parent e7505b2 commit 548eccc
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 106 deletions.
2 changes: 1 addition & 1 deletion src/gallia/cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -292,7 +292,7 @@ def build_cli(parsers: dict[str, Any], config: dict[str, Any]) -> None:
epilog=cls.EPILOG,
)
scanner = cls(subparser, config)
subparser.set_defaults(run_func=scanner.run)
subparser.set_defaults(run_func=scanner.entry_point)


def cmd_show_config(
Expand Down
229 changes: 124 additions & 105 deletions src/gallia/command.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,18 +13,18 @@
from abc import ABC, abstractmethod
from argparse import ArgumentParser, Namespace
from asyncio import Task
from copy import deepcopy
from datetime import datetime, timezone
from enum import Enum, IntEnum, unique
from importlib.metadata import EntryPoint, entry_points, version
from importlib.metadata import EntryPoint, entry_points
from pathlib import Path
from tempfile import gettempdir
from typing import Any, Optional, cast

import aiofiles
import msgspec

from gallia.db.db_handler import DBHandler
from gallia.log import get_logger, setup_logging
from gallia.log import get_logger, setup_logging, tz
from gallia.penlab import Dumpcap, PowerSupply, PowerSupplyURI
from gallia.transports.base import BaseTransport, TargetURI
from gallia.transports.can import ISOTPTransport, RawCANTransport
Expand All @@ -40,14 +40,28 @@
class ExitCodes(IntEnum):
SUCCESS = 0
GENERIC_ERROR = 1
SETUP_FAILED = 10
TEARDOWN_FAILED = 11
UNHANDLED_EXCEPTION = 2


@unique
class FileNames(Enum):
PROPERTIES_PRE = "PROPERTIES_PRE.json"
PROPERTIES_POST = "PROPERTIES_POST.json"
META = "META.json"


class CommandMeta(msgspec.Struct):
category: str
subcategory: Optional[str]
command: str


class RunMeta(msgspec.Struct):
command: list[str]
command_meta: CommandMeta
start_time: str
end_time: str
exit_code: int


def load_transport(target: TargetURI) -> BaseTransport:
Expand Down Expand Up @@ -114,9 +128,25 @@ def __init__(self, parser: ArgumentParser, config: dict[str, Any]) -> None:
self.logger = get_logger("gallia")
self.parser = parser
self.config = config
self.artifacts_dir = Path(".")
self.run_meta = RunMeta(
command=sys.argv,
command_meta=CommandMeta(
command=self.COMMAND,
category=self.CATEGORY,
subcategory=self.SUBCATEGORY,
),
start_time=datetime.now(tz).isoformat(),
exit_code=0,
end_time=0,
)
self.add_class_parser()
self.add_parser()

@abstractmethod
def run(self, args: Namespace) -> int:
...

def get_config_value(
self,
key: str,
Expand Down Expand Up @@ -226,9 +256,45 @@ def prepare_artifactsdir(

raise ValueError("base_dir or force_path must be different from None")

@abstractmethod
def run(self, args: Namespace) -> int:
...
def entry_point(self, args: Namespace) -> int:
if self.ARTIFACTSDIR:
self.artifacts_dir = self.prepare_artifactsdir(
args.artifacts_base,
args.artifacts_dir,
)
setup_logging(
self.get_log_level(args),
self.get_file_log_level(args),
self.artifacts_dir.joinpath("log.json.zst"),
)
else:
setup_logging(self.get_log_level(args))

exit_code = 0
try:
exit_code = self.run(args)
except (BrokenPipeError, ConnectionRefusedError) as e:
exit_code = 1
self.logger.critical(g_repr(e))
except KeyboardInterrupt:
exit_code = 128 + signal.SIGINT
# Ensure thet META.json gets written in the case a
# command calls sys.exit().
except SystemExit as e:
exit_code = e.code
except Exception:
exit_code = ExitCodes.UNHANDLED_EXCEPTION
traceback.print_exc()
else:
self.logger.info(f"Stored artifacts at {self.artifacts_dir}")
finally:
if self.ARTIFACTSDIR:
self.run_meta.exit_code = exit_code
self.run_meta.end_time = datetime.now(tz).isoformat()
data = msgspec.json.encode(self.run_meta)
self.artifacts_dir.joinpath(FileNames.META.value).write_bytes(data)

return exit_code


class Script(BaseCommand, ABC):
Expand All @@ -244,13 +310,8 @@ def main(self, args: Namespace) -> None:
...

def run(self, args: Namespace) -> int:
setup_logging(self.get_log_level(args))

try:
self.main(args)
return 0
except KeyboardInterrupt:
return 128 + signal.SIGINT
self.main(args)
return ExitCodes.SUCCESS


class AsyncScript(BaseCommand, ABC):
Expand All @@ -266,13 +327,8 @@ async def main(self, args: Namespace) -> None:
...

def run(self, args: Namespace) -> int:
setup_logging(self.get_log_level(args))

try:
asyncio.run(self.main(args))
return 0
except KeyboardInterrupt:
return 128 + signal.SIGINT
asyncio.run(self.main(args))
return ExitCodes.SUCCESS


class Scanner(BaseCommand, ABC):
Expand All @@ -297,7 +353,6 @@ class Scanner(BaseCommand, ABC):

def __init__(self, parser: ArgumentParser, config: dict[str, Any]) -> None:
super().__init__(parser, config)
self.artifacts_dir: Path
self.db_handler: Optional[DBHandler] = None
self.power_supply: Optional[PowerSupply] = None
self.dumpcap: Optional[Dumpcap] = None
Expand All @@ -306,6 +361,38 @@ def __init__(self, parser: ArgumentParser, config: dict[str, Any]) -> None:
async def main(self, args: Namespace) -> None:
...

async def _db_insert_run_meta(self, args: Namespace) -> None:
if args.db is not None:
self.db_handler = DBHandler(args.db)
await self.db_handler.connect()

await self.db_handler.insert_run_meta(
script=sys.argv[0].split()[-1],
arguments=sys.argv[1:],
start_time=datetime.now(timezone.utc).astimezone(),
path=self.artifacts_dir,
)

async def _db_finish_run_meta(self, args: Namespace, exit_code: int) -> None:
if self.db_handler is not None and self.db_handler.connection is not None:
if self.db_handler.meta is not None:
try:
await self.db_handler.complete_run_meta(
datetime.now(timezone.utc).astimezone(),
exit_code,
)
except Exception as e:
self.logger.warning(
f"Could not write the run meta to the database: {g_repr(e)}"
)

try:
await self.db_handler.disconnect()
except Exception as e:
self.logger.error(
f"Could not close the database connection properly: {g_repr(e)}"
)

async def setup(self, args: Namespace) -> None:
if args.target is None:
self.parser.error("--target is required")
Expand Down Expand Up @@ -373,92 +460,22 @@ def add_class_parser(self) -> None:
),
)

async def _run(self, args: Namespace) -> int:
exit_code: int = ExitCodes.SUCCESS

async def _run(self, args: Namespace) -> None:
await self.setup(args)
try:
if args.db is not None:
self.db_handler = DBHandler(args.db)
await self.db_handler.connect()

await self.db_handler.insert_run_meta(
script=sys.argv[0].split()[-1],
arguments=sys.argv[1:],
start_time=datetime.now(timezone.utc).astimezone(),
path=self.artifacts_dir,
)

try:
await self.setup(args)
except BrokenPipeError as e:
exit_code = ExitCodes.GENERIC_ERROR
self.logger.critical(g_repr(e))
except Exception as e:
self.logger.exception(f"setup failed: {g_repr(e)}")
sys.exit(ExitCodes.SETUP_FAILED)

try:
try:
await self.main(args)
except Exception as e:
exit_code = ExitCodes.GENERIC_ERROR
self.logger.critical(g_repr(e))
traceback.print_exc()
finally:
try:
await self.teardown(args)
except Exception as e:
self.logger.critical(f"teardown failed: {g_repr(e)}")
sys.exit(ExitCodes.TEARDOWN_FAILED)
return exit_code
except KeyboardInterrupt:
exit_code = 128 + signal.SIGINT
raise
except SystemExit as se:
exit_code = se.code
raise
await self.main(args)
finally:
if self.db_handler is not None and self.db_handler.connection is not None:
if self.db_handler.meta is not None:
try:
await self.db_handler.complete_run_meta(
datetime.now(timezone.utc).astimezone(), exit_code
)
except Exception as e:
self.logger.warning(
f"Could not write the run meta to the database: {g_repr(e)}"
)

try:
await self.db_handler.disconnect()
except Exception as e:
self.logger.error(
f"Could not close the database connection properly: {g_repr(e)}"
)
await self.teardown(args)

def run(self, args: Namespace) -> int:
self.artifacts_dir = self.prepare_artifactsdir(
args.artifacts_base,
args.artifacts_dir,
)

setup_logging(
self.get_log_level(args),
self.get_file_log_level(args),
self.artifacts_dir.joinpath("log.json.zst"),
)
asyncio.run(self._run(args))
return ExitCodes.SUCCESS

argv = deepcopy(sys.argv)
argv[0] = Path(sys.argv[0]).name
self.logger.info(
f'Starting "{sys.argv[0]}" ({version("gallia")}) with [{" ".join(argv)}]'
)
self.logger.info(f"Storing artifacts at {self.artifacts_dir}")

try:
return asyncio.run(self._run(args))
except KeyboardInterrupt:
return 128 + signal.SIGINT
def entry_point(self, args: Namespace) -> int:
asyncio.run(self._db_insert_run_meta(args))
exit_code = super().entry_point(args)
asyncio.run(self._db_finish_run_meta(args, exit_code))
return exit_code


class UDSScanner(Scanner):
Expand Down Expand Up @@ -617,7 +634,9 @@ async def setup(self, args: Namespace) -> None:

if self.db_handler is not None:
try:
await self.db_handler.insert_scan_run(str(args.target))
# No idea, but str(args.target) fails with a strange traceback.
# Lets use the attribute directly…
await self.db_handler.insert_scan_run(args.target.raw)
self._apply_implicit_logging_setting()
except Exception as e:
self.logger.warning(
Expand Down

0 comments on commit 548eccc

Please sign in to comment.