|
Server : Apache System : Linux srv.rainic.com 4.18.0-553.47.1.el8_10.x86_64 #1 SMP Wed Apr 2 05:45:37 EDT 2025 x86_64 User : rainic ( 1014) PHP Version : 7.4.33 Disable Function : exec,passthru,shell_exec,system Directory : /opt/imunify360/venv/lib/python3.11/site-packages/defence360agent/utils/ |
Upload File : |
import functools
import logging
import pickle
from asyncio import iscoroutinefunction
from typing import Union, Callable
logger = logging.getLogger(__name__)
def _dump(path, obj):
# separate function to mock in tests
with open(path, "wb") as w:
pickle.dump(obj, w)
def serialize_attr(*, path: str, attr: str):
"""
Make decorator to serialize an object or object's attribute
:param path: path to file to serialize into
:param attr: attribute name to serialize
"""
def decorator(f):
@functools.wraps(f)
def wrapper(self, *args, **kwargs):
result = f(self, *args, **kwargs)
obj = getattr(self, attr)
logger.debug("Write %r to %r", obj, path)
_dump(path, obj)
return result
@functools.wraps(f)
async def async_wrapper(self, *args, **kwargs):
result = await f(self, *args, **kwargs)
obj = getattr(self, attr)
logger.debug("Write %r to %r", obj, path)
_dump(path, obj)
return result
if iscoroutinefunction(f):
return async_wrapper
return wrapper
return decorator
def unserialize(*, path: str, fallback: Union[Callable, object] = None):
"""
Restore an object from file
:param path: path to read from
:param fallback: return or call it if unserialization fails
:return:
"""
try:
with open(path, "rb") as r:
return pickle.load(r)
except FileNotFoundError:
logger.warning("Can't find %s to unserialize", path)
except Exception as e:
logger.error("Unserialize failed with %r. Returning an fallback", e)
return fallback() if callable(fallback) else fallback