관리-도구

편집 파일: revisium.py

"""
This program is free software: you can redistribute it and/or modify it under
the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License,
or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
See the GNU General Public License for more details.

You should have received a copy of the GNU General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.

This software is also available under ImunifyAV commercial license,
see <https://www.imunify360.com/legal/eula>
"""
import csv
import json
import os
import sys
import time
from contextlib import suppress
from pathlib import Path
from typing import Optional

from defence360agent.contracts.config import Core
from defence360agent.utils import snake_case

class RevisiumTempFile:
    ext = "txt"

def __init__(self, tempdir=None, mode=None):
        self.tempdir = Path(Core.TMPDIR) if tempdir is None else Path(tempdir)
        self.mode = mode

def __enter__(self):
        filename = "%s_%d.%s" % (self.name, int(time.time() * 10e6), self.ext)
        self._path = self.tempdir / filename
        if self.mode:
            self._path.touch()
            os.chmod(str(self._path), mode=self.mode)
        return self

def __exit__(self, *_):
        with suppress(FileNotFoundError):
            self._path.unlink()

@property
    def filename(self):
        return str(self._path)

@property
    def name(self):
        return snake_case(self.__class__.__name__)

class RevisiumJsonFile(RevisiumTempFile):
    ext = "json"

def read(self):
        with self._path.open() as f:
            return json.load(f)

class RevisiumCSVFile(RevisiumTempFile):
    ext = "csv"

@staticmethod
    def _convert_boolean(value: Optional[str]):
        if value:
            low_value = value.lower()
            if low_value in ["true", "false"]:
                return low_value == "true"
        elif value is None:
            return ""
        return value

def read(self):
        """
        Reads a CSV file and returns its contents as a list of dictionaries.

The file must be read in binary mode and decoded manually to preserve the original line ending and carriage
        return characters.
        """
        with self._path.open("rb") as f:
            contents = [
                line.decode(
                    sys.getfilesystemencoding(),
                    errors=sys.getfilesystemencodeerrors(),
                )
                for line in f
            ]
            reader = csv.DictReader(
                contents, fieldnames=["f", "e", "d", "s", "t", "r"]
            )
            return [
                {
                    field: self._convert_boolean(value)
                    for field, value in row.items()
                    if field is not None  # ignore unknown fields
                }
                for row in reader
            ]