MAINHACK

Viewing / Editing File: utils.py

File is not writable. Editing disabled.
import pickle
import asyncio
import socket
from functools import lru_cache, wraps
from itertools import chain
from pathlib import Path
from typing import Optional, Tuple, Generator

import yaml

from defence360agent.contracts.config import SimpleRpc as Config
from defence360agent.model.simplification import run_in_executor
from defence360agent.rpc_tools.validate import ValidationError
from defence360agent.utils import antivirus_mode, is_root_user


def is_running():
    """Check if Agent (or Sensor) instance is running"""
    if Config.SOCKET_ACTIVATION:
        # TODO: check that process is exist already
        # As systemd use sockets in antivirus mode
        # we could not use it for check
        return False

    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock:
        try:
            if is_root_user():
                sock.connect(Config.SOCKET_PATH)
            else:
                sock.connect(Config.NON_ROOT_SOCKET_PATH)
        except (ConnectionRefusedError, FileNotFoundError):
            return False
        else:
            return True


def _find_schema(base=None, schema_dir="schema"):
    for path in find_schema_files(base, schema_dir):
        if (p := path.with_suffix(".pickle")).exists():
            document = pickle.loads(p.read_bytes())
        else:
            document = yaml.safe_load(path.read_text())

        for k, v in document.items():
            # converting keys - from strings to tuples
            yield tuple(k.split(" ")), v


def find_schema_files(
    base=None, schema_dir="schema"
) -> Generator[Path, None, None]:
    p = base / schema_dir
    for path in [*p.rglob("*.yaml"), *p.rglob("*.yml")]:
        yield path


def get_schema_paths(paths: Optional[Tuple[Path]] = None):
    result = list(paths)[:] if paths else []
    path = Path(__file__).parent.parent / "simple_rpc"
    result.extend(
        [
            path,
            path.parent / "feature_management" / "rpc",
        ]
    )
    return result


@lru_cache(1)
def prepare_schema(paths):
    schema = []
    for base_path in get_schema_paths(paths):
        schema.append(_find_schema(base_path))
    return dict(chain(*schema))


def run_in_executor_decorator(f):
    @wraps(f)
    async def wrapper(*args, **kwargs):
        return await run_in_executor(
            asyncio.get_event_loop(), lambda: f(*args, **kwargs)
        )

    return wrapper


def generate_warnings(
    affected,
    not_affected,
    dest_listname,
    all_list,
    success_warning,
    failure_warning,
    in_another_list_warning=None,
):
    """
    :param list affected: IPs that were changed during operation
    :param list of tuples || list of str not_affected: IPs & it's listnames
            that weren't changed during operation
    :param list all_list: list of all IPs that take place in operation
    :param str success_warning: msg if IP was changed
    :param str failure_warning: msg if IPs wasn't changed and it's absent
            in any other lists
    :param str in_another_list_warning: msg if IPs wasn't changed , however
            it present in another list
    :return list of st warnings: msg to be printed
    """
    warnings = []
    for item in not_affected:
        record, listname = item["rec"], item.get("listname", dest_listname)
        if listname and in_another_list_warning:
            warnings.append(in_another_list_warning.format(record, listname))
        else:
            warnings.append(failure_warning.format(record, dest_listname))

    num_deleted = len(affected)
    total_num = len(all_list)

    if total_num > 1 and total_num != num_deleted:
        warnings.append(success_warning.format(num_deleted, total_num))

    if warnings:
        raise ValidationError(warnings)

    return {}
Cancel / Back
[ MAINHACK ]

Viewing / Editing File: utils.py

Server Info

System Features

User Info
