recurvedata-lib 0.1.487__py2.py3-none-any.whl

recurvedata/utils/attrdict.py

@@ -0,0 +1,19 @@
+from typing import Any, KeysView, Union
+
+
+class AttrDict(dict):
+    """A dict that allows for attribute-style access."""
+
+    def __getattr__(self, item: str) -> Union[Any, "AttrDict"]:
+        if item not in self:
+            return None
+        value = self[item]
+        if isinstance(value, dict):
+            return AttrDict(value)
+        return value
+
+    def __setattr__(self, key: str, value: Any) -> None:
+        self[key] = value
+
+    def __dir__(self) -> KeysView[str]:
+        return self.keys()

recurvedata/utils/command_helper.py

@@ -0,0 +1,20 @@
+import contextlib
+import json
+import os
+import traceback
+
+
+@contextlib.contextmanager
+def capture_error_to_file(filename: str, suppress_error: bool = False):
+    if not filename:
+        yield
+    else:
+        try:
+            yield
+        except Exception as e:
+            error_info = {"traceback": traceback.format_exc(), "error_message": str(e), "is_success": False}
+            os.makedirs(os.path.dirname(filename), exist_ok=True)
+            with open(filename, "w") as f:
+                json.dump(error_info, f, indent=2)
+            if not suppress_error:
+                raise

recurvedata/utils/compat.py

@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+import sys
+
+PY38 = sys.version_info >= (3, 8)
+PY39 = sys.version_info >= (3, 9)
+PY310 = sys.version_info >= (3, 10)
+PY311 = sys.version_info >= (3, 11)
+PY312 = sys.version_info >= (3, 12)
+
+is_osx = sys.platform == "darwin"
+is_win = sys.platform == "win32"

recurvedata/utils/compression.py

@@ -0,0 +1,203 @@
+import bz2
+import gzip
+import logging
+import os
+import shutil
+import struct
+import tarfile
+import tempfile
+import zipfile
+import zlib
+from typing import Callable, NamedTuple, Optional
+
+from recurvedata.utils import files, shell
+
+logger = logging.getLogger(__name__)
+
+
+class _Config(NamedTuple):
+    opener: Callable
+    ext: str
+    compress_cmd: str
+    decompress_cmd: str
+
+
+_gzip_cfg = _Config(gzip.open, "gz", "gzip", "gzip -d")
+_bzip2_cfg = _Config(bz2.open, "bz2", "bzip2", "bzip2 -d")
+
+
+def gzip_compress(src_file: str, dst_file: str = None, using_cmd: bool = False, inplace: bool = False) -> str:
+    """Compress a file using gzip
+
+    Args:
+        src_file: the path of input file
+        dst_file: the path of output file, a temporary filename will be made otherwise
+        using_cmd: use the gzip command line instead of Python gzip to speedup
+        inplace: replace the src_file with dst_file if True
+
+    Returns:
+        the dst_file if inplace is False, otherwise the src_file
+    """
+    return _compress_file(_gzip_cfg, src_file, dst_file, using_cmd, inplace)
+
+
+def gzip_decompress(src_file: str, dst_file: str = None, using_cmd: bool = False, inplace: bool = False) -> str:
+    """Decompress a .gz file
+
+    Args:
+        src_file: the path of input file
+        dst_file: the path of output file, a temporary filename will be made otherwise
+        using_cmd: use the gzip command line instead of Python gzip to speedup
+        inplace: replace the filename with target_filename if True
+
+    Returns:
+        the dst_file if inplace is False, otherwise the src_file
+    """
+    return _decompress_file(_gzip_cfg, src_file, dst_file, using_cmd, inplace)
+
+
+def bzip2_compress(src_file: str, dst_file: str = None, using_cmd: bool = False, inplace: bool = False) -> str:
+    """Compress a file using bzip2
+
+    Args:
+        src_file: the path of input file
+        dst_file: the path of output file, a temporary filename will be made otherwise
+        using_cmd: use the bzip2 command line instead of Python bzip2 to speedup
+        inplace: replace the filename with target_filename if True
+
+    Returns:
+        the dst_file if inplace is False, otherwise the src_file
+    """
+    return _compress_file(_bzip2_cfg, src_file, dst_file, using_cmd, inplace)
+
+
+def bzip2_decompress(src_file: str, dst_file: str = None, using_cmd: bool = False) -> str:
+    """Decompress a .bz2 file
+
+    Args:
+        src_file: the path of input file
+        dst_file: the path of output file, a temporary filename will be made otherwise
+        using_cmd: use the bzip2 command line instead of Python bzip2 to speedup
+
+    Returns:
+        the dst_file if inplace is False, otherwise the src_file
+    """
+    return _decompress_file(_bzip2_cfg, src_file, dst_file, using_cmd, False)
+
+
+def _compress_file(
+    cfg: _Config, src_file: str, dst_file: str = None, using_cmd: bool = False, inplace: bool = False
+) -> str:
+    if dst_file is None:
+        dst_file = files.new_tempfile(suffix=f".{cfg.ext}")
+
+    if using_cmd:
+        shell.run(f"{cfg.compress_cmd} {src_file} -c > {dst_file}", logger)
+    else:
+        with open(src_file, "rb") as f_in, cfg.opener(dst_file, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+    return files.replace_file_with_temp(dst_file, src_file, inplace)
+
+
+def _decompress_file(
+    cfg: _Config, src_file: str, dst_file: str = None, using_cmd: bool = False, inplace: bool = False
+) -> str:
+    if dst_file is None:
+        dst_file = files.new_tempfile()
+
+    if using_cmd:
+        shell.run(f"{cfg.decompress_cmd} {src_file} -c > {dst_file}", logger)
+    else:
+        with cfg.opener(src_file, "rb") as f_in, open(dst_file, "wb") as f_out:
+            shutil.copyfileobj(f_in, f_out)
+
+    return files.replace_file_with_temp(dst_file, src_file, inplace)
+
+
+def zip_compress(src_file: str, dst_file: str = None, using_cmd: bool = False, arcname: str = None) -> str:
+    """Compress a file using zip
+
+    Args:
+        src_file: the path of input file
+        dst_file: the path of output file, a temporary filename will be made otherwise
+        using_cmd: use the zip command line instead of Python ZipFile to speedup
+        arcname: filename in the archive file, only supported with using_cmd=False
+
+    Returns:
+        the target_filename
+    """
+    if dst_file is None:
+        dst_file = files.new_tempfile(suffix=".zip")
+
+    directory, basename = os.path.split(src_file.rstrip("/"))
+
+    if using_cmd:
+        # 先删除生成的临时文件，只使用生成的文件名，要不然会报错
+        # zip warning: missing end signature--probably not a zip file (did you
+        # zip warning: remember to use binary mode when you transferred it?)
+        # zip warning: (if you are trying to read a damaged archive try -F)
+        files.remove_files_safely(dst_file)
+        if arcname is not None:
+            logger.warning("arcname is not supported while using cmd")
+        shell.run(f"cd {directory} && zip -r {dst_file} {basename}", logger)
+        return dst_file
+
+    with zipfile.ZipFile(dst_file, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        zf.write(src_file, arcname=arcname or basename)
+    return dst_file
+
+
+def zip_decompress(src_file: str, target_directory: str = None, using_cmd: bool = False) -> str:
+    """Decompress a .zip file into a directory
+
+    Args:
+        src_file: the path of input file
+        target_directory: the path of output directory, a temporary directory will be made otherwise
+        using_cmd: use the unzip command line instead of Python ZipFile to speedup
+
+    Returns:
+        the output directory
+    """
+    if not target_directory:
+        target_directory = tempfile.mkdtemp()
+
+    if using_cmd:
+        shell.run(f"unzip {src_file} -d {target_directory}", logger)
+        return target_directory
+
+    with zipfile.ZipFile(src_file, "r") as zf:
+        zf.extractall(target_directory)
+    return target_directory
+
+
+def mysql_compress(value: bytes) -> Optional[bytes]:
+    """A Python implementation of COMPRESS function of MySQL."""
+    if value is None:
+        return None
+    if value == b"":
+        return b""
+    size: bytes = struct.pack("I", len(value))
+    data: bytes = zlib.compress(value)
+    return size + data
+
+
+def mysql_uncompress(value: bytes) -> bytes:
+    """A Python implementation of UNCOMPRESS function of MySQL.
+
+    Used to decompress result of COMPRESS function.
+
+    https://dev.mysql.com/doc/refman/5.7/en/encryption-functions.html#function_compress
+    """
+    size_byte: int = 4
+    if not value or len(value) < size_byte:
+        return value
+
+    return zlib.decompress(value[size_byte:])
+
+
+def tar_gzip_uncompress(tar_gz_path, extract_path):
+    logger.info(f"extract tar.gz {tar_gz_path} to {extract_path}")
+    os.makedirs(extract_path, exist_ok=True)
+    with tarfile.open(tar_gz_path, "r:gz") as tar:
+        tar.extractall(path=extract_path)

recurvedata/utils/crontab.py

@@ -0,0 +1,42 @@
+import datetime
+from typing import Optional
+
+import croniter
+
+from recurvedata.utils.date_time import DATELIKE
+
+
+def normalize_schedule_interval(schedule_interval: str) -> str:
+    mapping = {
+        "@once": None,
+        "@hourly": "0 * * * *",
+        "@daily": "0 0 * * *",
+        "@weekly": "0 0 * * 0",
+        "@monthly": "0 0 1 * *",
+        "@yearly": "0 0 1 1 *",
+    }
+    return mapping.get(schedule_interval, schedule_interval)
+
+
+def _get_schedule(schedule_interval: str, dttm: DATELIKE, is_next: bool = False) -> Optional[DATELIKE]:
+    cron_spec = normalize_schedule_interval(schedule_interval)
+    if not cron_spec:
+        return None
+    start_time = dttm
+    cron = croniter.croniter(cron_spec, start_time)
+    if is_next:
+        value = cron.get_next(datetime.datetime)
+    else:
+        value = cron.get_prev(datetime.datetime)
+    return value
+
+
+def next_schedule(schedule_interval: str, dttm: DATELIKE) -> DATELIKE:
+    return _get_schedule(schedule_interval, dttm, is_next=True)
+
+
+def previous_schedule(schedule_interval: str, dttm: DATELIKE) -> DATELIKE:
+    return _get_schedule(schedule_interval, dttm, is_next=False)
+
+
+get_schedule = _get_schedule

recurvedata/utils/crypto_util.py

@@ -0,0 +1,305 @@
+import base64
+import hashlib
+import json
+import logging
+import os
+from enum import Enum
+from functools import lru_cache
+from typing import Literal, NamedTuple
+
+from recurvedata.client.client import Client
+from recurvedata.utils.imports import MockModule
+from recurvedata.utils.registry import register_func
+
+try:
+    from cryptography.hazmat.primitives import hashes, padding, serialization
+    from cryptography.hazmat.primitives.asymmetric import padding as asymmetric_padding
+    from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+except ImportError:
+    mock_module = MockModule("cryptography")
+    hashes = mock_module
+    padding = mock_module
+    serialization = mock_module
+    asymmetric_padding = mock_module
+    Cipher = mock_module
+    algorithms = mock_module
+    modes = mock_module
+
+logger = logging.getLogger(__name__)
+
+
+class CryptoMethod(str, Enum):
+    AES_128 = "AES-128"
+    AES_256 = "AES-256"
+    RSA_2048 = "RSA-2048"
+    RSA_4096 = "RSA-4096"
+
+
+class RSAKeyPair(NamedTuple):
+    public_key: bytes
+    private_key: bytes
+
+
+class CryptoUtil:
+    """Utility class for encryption operations"""
+
+    @staticmethod
+    @lru_cache
+    def fetch_key_data(key_name: str) -> tuple[str, str]:
+        """Fetch key data from server using key name"""
+        client = Client()
+        response = client.request("GET", f"/api/executor/keys/{key_name}")
+        key_data = response["key_data"]
+        encryption_method = response["encryption_method"]
+        return key_data, encryption_method
+
+    @staticmethod
+    @lru_cache
+    def get_key_data(key_name: str) -> bytes | RSAKeyPair:
+        """Get and process key data from server using key name
+
+        Args:
+            key_name: Name of the key to fetch
+
+        Returns:
+            Processed key data based on encryption method from server
+        """
+        key_data, encryption_method = CryptoUtil.fetch_key_data(key_name)
+        return CryptoUtil._decode_key_data(key_data, encryption_method)
+
+    @staticmethod
+    def _decode_key_data(key_data: str, encryption_method: str) -> bytes | RSAKeyPair:
+        if encryption_method in (CryptoMethod.AES_128, CryptoMethod.AES_256):
+            try:
+                key = base64.urlsafe_b64decode(key_data)
+            except Exception as e:
+                raise ValueError("Invalid base64 encoding") from e
+            if len(key) not in (16, 24, 32):
+                raise ValueError("AES key must be 16, 24, or 32 bytes long")
+            return key
+        elif encryption_method in (CryptoMethod.RSA_2048, CryptoMethod.RSA_4096):
+            try:
+                key_dict = json.loads(key_data)
+            except Exception as e:
+                raise ValueError("Invalid RSA key data") from e
+            if not isinstance(key_dict, dict):
+                raise ValueError("Invalid RSA key data")
+            if "public_key" not in key_dict or "private_key" not in key_dict:
+                raise ValueError("Public key or private key is missing")
+            return RSAKeyPair(
+                public_key=key_dict["public_key"].encode("utf-8"),
+                private_key=key_dict["private_key"].encode("utf-8"),
+            )
+        else:
+            raise ValueError(f"Unsupported encryption method: {encryption_method}")
+
+    @staticmethod
+    def aes_encrypt_with_key(
+        key: bytes, data: str | bytes | None, mode: Literal["ECB", "CBC"] = "ECB", iv: str | bytes | None = None
+    ) -> bytes | None:
+        """Encrypt data using AES with provided key
+
+        Args:
+            key: AES key bytes
+            data: Data to encrypt
+            mode: AES mode ('ECB' or 'CBC')
+            iv: Initialization vector for CBC mode
+
+        Returns:
+            Encrypted bytes (with IV prepended for CBC mode)
+        """
+        if data is None:
+            return None
+
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+
+        padder = padding.PKCS7(algorithms.AES.block_size).padder()
+        padded_data = padder.update(data) + padder.finalize()
+
+        if mode.upper() == "ECB":
+            cipher = Cipher(algorithms.AES(key), modes.ECB()).encryptor()
+            return cipher.update(padded_data) + cipher.finalize()
+        elif mode.upper() == "CBC":
+            if iv is None:
+                processed_iv = os.urandom(16)
+            else:
+                if isinstance(iv, str):
+                    iv = iv.encode("utf-8")
+                # Hash the IV to get exactly 16 bytes
+                processed_iv = hashlib.md5(iv).digest()
+
+            cipher = Cipher(algorithms.AES(key), modes.CBC(processed_iv)).encryptor()
+            return processed_iv + cipher.update(padded_data) + cipher.finalize()
+        else:
+            raise ValueError(f"Unsupported AES mode: {mode}")
+
+    @staticmethod
+    def aes_decrypt_with_key(key: bytes, data: bytes | None, mode: Literal["ECB", "CBC"] = "ECB") -> bytes | None:
+        """Decrypt data using AES with provided key"""
+        if data is None:
+            return None
+
+        if mode.upper() == "ECB":
+            cipher = Cipher(algorithms.AES(key), modes.ECB()).decryptor()
+            padded_data = cipher.update(data) + cipher.finalize()
+        elif mode.upper() == "CBC":
+            iv = data[:16]
+            cipher = Cipher(algorithms.AES(key), modes.CBC(iv)).decryptor()
+            padded_data = cipher.update(data[16:]) + cipher.finalize()
+        else:
+            raise ValueError(f"Unsupported AES mode: {mode}")
+
+        unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
+        return unpadder.update(padded_data) + unpadder.finalize()
+
+    @staticmethod
+    def aes_encrypt(
+        key_name: str, data: str | bytes | None, mode: Literal["ECB", "CBC"] = "ECB", iv: str | bytes | None = None
+    ) -> bytes | None:
+        """Encrypt data using AES with key from key store"""
+        key: bytes = CryptoUtil.get_key_data(key_name)
+        return CryptoUtil.aes_encrypt_with_key(key, data, mode, iv)
+
+    @staticmethod
+    def aes_decrypt(key_name: str, data: bytes | None, mode: Literal["ECB", "CBC"] = "ECB") -> bytes | None:
+        """Decrypt data using AES with key from key store"""
+        key: bytes = CryptoUtil.get_key_data(key_name)
+        return CryptoUtil.aes_decrypt_with_key(key, data, mode)
+
+    @staticmethod
+    def rsa_encrypt_with_key(public_key: bytes, data: str | bytes) -> bytes:
+        """Encrypt data using RSA with provided public key"""
+        public_key_obj = serialization.load_pem_public_key(public_key)
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+
+        return public_key_obj.encrypt(
+            data,
+            asymmetric_padding.OAEP(
+                mgf=asymmetric_padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+
+    @staticmethod
+    def rsa_decrypt_with_key(private_key: bytes, data: bytes) -> bytes:
+        """Decrypt data using RSA with provided private key"""
+        private_key_obj = serialization.load_pem_private_key(private_key, password=None)
+        return private_key_obj.decrypt(
+            data,
+            asymmetric_padding.OAEP(
+                mgf=asymmetric_padding.MGF1(algorithm=hashes.SHA256()),
+                algorithm=hashes.SHA256(),
+                label=None,
+            ),
+        )
+
+    @staticmethod
+    def rsa_encrypt(key_name: str, data: str | bytes | None) -> bytes | None:
+        """Encrypt data using RSA with key from key store"""
+        if data is None:
+            return None
+        key_pair: RSAKeyPair = CryptoUtil.get_key_data(key_name)
+        return CryptoUtil.rsa_encrypt_with_key(key_pair.public_key, data)
+
+    @staticmethod
+    def rsa_decrypt(key_name: str, data: bytes | None) -> bytes | None:
+        """Decrypt data using RSA with key from key store"""
+        if data is None:
+            return None
+        key_pair: RSAKeyPair = CryptoUtil.get_key_data(key_name)
+        return CryptoUtil.rsa_decrypt_with_key(key_pair.private_key, data)
+
+    @staticmethod
+    def base64_encode(data: bytes | None) -> str | None:
+        """Convert bytes data to base64 encoded string.
+
+        Args:
+            data: The bytes data to encode. Can be None.
+
+        Returns:
+            Base64 encoded string if input is not None, otherwise None.
+        """
+        if data is None:
+            return None
+        return base64.b64encode(data).decode("utf-8")
+
+    @staticmethod
+    def base64_decode(data: str | None) -> bytes | None:
+        """Convert base64 encoded encrypted data back to bytes
+
+        Args:
+            data: Base64 encoded encrypted data string. Can be None.
+
+        Returns:
+            Decoded bytes data if input is not None, otherwise None.
+        """
+        if data is None:
+            return None
+        return base64.b64decode(data.encode("utf-8"))
+
+    @staticmethod
+    def md5(data: str | bytes | None) -> str | None:
+        """Calculate MD5 hash of data
+
+        Args:
+            data: Input data as string or bytes
+
+        Returns:
+            MD5 hash as hex string
+        """
+        return CryptoUtil._hash(data, "md5")
+
+    @staticmethod
+    def sha1(data: str | bytes | None) -> str | None:
+        """Calculate SHA1 hash of data
+
+        Args:
+            data: Input data as string or bytes
+
+        Returns:
+            SHA1 hash as hex string
+        """
+        return CryptoUtil._hash(data, "sha1")
+
+    @staticmethod
+    def sha256(data: str | bytes | None) -> str | None:
+        """Calculate SHA256 hash of data
+
+        Args:
+            data: Input data as string or bytes
+
+        Returns:
+            SHA256 hash as hex string
+        """
+        return CryptoUtil._hash(data, "sha256")
+
+    @staticmethod
+    def _hash(data: str | bytes | None, algorithm: Literal["md5", "sha1", "sha256"]) -> str | None:
+        if data is None:
+            return None
+        if isinstance(data, str):
+            data = data.encode("utf-8")
+        return hashlib.new(algorithm, data).hexdigest()
+
+
+@register_func
+def secret(key_name: str) -> str:
+    """Get secret from server using key name"""
+    try:
+        key_data, _ = CryptoUtil.fetch_key_data(key_name)
+        return key_data
+    except Exception as e:
+        logger.warning(f"Failed to get secret for key {key_name}: {e}")
+
+        return key_name
+
+
+aes_encrypt = CryptoUtil.aes_encrypt
+aes_decrypt = CryptoUtil.aes_decrypt
+rsa_encrypt = CryptoUtil.rsa_encrypt
+rsa_decrypt = CryptoUtil.rsa_decrypt
+base64_encode = CryptoUtil.base64_encode
+base64_decode = CryptoUtil.base64_decode

recurvedata/utils/dataclass.py

@@ -0,0 +1,11 @@
+from dataclasses import fields
+from typing import Any
+
+from recurvedata.utils.helpers import extract_dict
+
+
+def init_dataclass_from_dict(dc_cls, arg_dct: dict[str, Any], **kwargs):
+    field_names = [f.name for f in fields(dc_cls) if f.init]
+
+    valid_arg_dct = extract_dict(arg_dct, field_names)
+    return dc_cls(**valid_arg_dct, **kwargs)