1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142
|
"""Dump tools for AsusRouter.
This module contains all needed to dump the raw data from the router
for the following analysis and testing.
"""
from __future__ import annotations
import asyncio
from datetime import UTC, datetime
import json
import os
from types import TracebackType
from typing import Any, Self
import zipfile
from asusrouter.modules.endpoint import Endpoint
class AsusRouterDump:
"""AsusRouter Dumper."""
_endpoint: Endpoint
_datetime: str
def __init__(
self, output_folder: str, full_dump: bool = False, archive: bool = True
) -> None:
"""Initialize."""
self.log: dict[str, str] = {}
self._output_folder = output_folder
self.full_dump = full_dump
self.zip = archive
self._init_datetime = datetime.now(UTC)
if self.zip:
self._zipfile = zipfile.ZipFile(
os.path.join( # noqa: PTH118
self._output_folder,
(
"AsusRouter-"
f"{self._init_datetime.isoformat().replace(':', '-')}"
".zip"
),
),
"w",
)
else:
# Create subfolder for the dump
self._output_folder = os.path.join( # noqa: PTH118
output_folder,
f"AsusRouter-"
f"{self._init_datetime.isoformat().replace(':', '-')}",
)
os.makedirs(self._output_folder) # noqa: PTH103
def __enter__(self) -> Self:
"""Enter the runtime context related to this object."""
return self
def __exit__(
self,
exc_type: type[BaseException] | None,
exc_value: BaseException | None,
traceback: TracebackType | None,
) -> None:
"""Exit the runtime context and save the log."""
if self.zip:
# Write the log data to the zip file
self._zipfile.writestr(
"log.json", json.dumps(self.log, default=str)
)
self._zipfile.close()
else:
log_filename = os.path.join(self._output_folder, "log.json") # noqa: PTH118
with open(log_filename, "w", encoding="utf-8") as f: # noqa: PTH123
json.dump(self.log, f, default=str)
async def dump(
self,
endpoint: Endpoint,
payload: dict[str, Any],
resp_status: int,
resp_headers: dict[str, Any],
resp_content: bytes,
) -> None:
"""Dump the data."""
self._endpoint = endpoint
self._datetime = datetime.now(UTC).isoformat().replace(":", "-")
self.log[self._datetime] = f"Endpoint.{endpoint.name}"
# Write the response content to a file
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._write_content, resp_content)
# Write the metadata to a file
metadata = {
"endpoint": endpoint,
"payload": payload,
"resp_status": resp_status,
"resp_headers": resp_headers,
}
await loop.run_in_executor(None, self._write_metadata, metadata)
def _write_content(self, content: str) -> None:
"""Write the content to a file."""
if self.full_dump:
filename = f"{self._datetime}-{self._endpoint}.content"
if self.zip:
self._zipfile.writestr(filename, content)
else:
with open( # noqa: PTH123
os.path.join(self._output_folder, filename), # noqa: PTH118
"w",
encoding="utf-8",
) as f:
f.write(content)
def _write_metadata(self, metadata: dict[str, Any]) -> None:
"""Write the metadata to a file."""
if self.full_dump:
filename = f"{self._datetime}-{self._endpoint}.json"
if self.zip:
self._zipfile.writestr(
filename, json.dumps(metadata, default=str)
)
else:
with open( # noqa: PTH123
os.path.join(self._output_folder, filename), # noqa: PTH118
"w",
encoding="utf-8",
) as f:
json.dump(metadata, f, default=str)
|