from __future__ import annotations
import ctypes
import os
from pathlib import Path
__all__ = [
"FfiError",
"LibraryNotFoundError",
"find_cdylib",
"sha256",
"sha1",
"sha512",
"hmac_sha1",
"hmac_sha256",
"hmac_sha512",
"xxh64",
"murmur3_x64_128",
"crc32c",
"base64_encode",
"base64_decode",
"xoshiro256_fill",
"crc32",
"adler32",
"fnv1a64",
"splitmix64_fill",
"STATUS_OK",
"STATUS_INVALID",
"STATUS_REJECTED",
]
STATUS_OK = 0
STATUS_INVALID = -1
STATUS_REJECTED = -2
CDYLIB_NAMES = ("pith_digest.dll", "libpith_digest.so", "libpith_digest.dylib")
class LibraryNotFoundError(OSError):
class FfiError(Exception):
def __init__(self, op: str, status: int) -> None:
kind = {
STATUS_INVALID: "invalid argument",
STATUS_REJECTED: "input rejected",
}.get(status, "unknown failure")
super().__init__(f"{op} failed: {kind} (status {status})")
self.status = status
def find_cdylib() -> Path:
explicit = os.environ.get("PITH_CDYLIB")
if explicit:
p = Path(explicit)
if p.is_file():
return p
env_dir = os.environ.get("PITH_CDYLIB_DIR")
candidates: list[Path] = []
if env_dir:
env_dir_path = Path(env_dir)
candidates.append(env_dir_path)
if not env_dir_path.is_absolute():
candidates.append(Path.cwd() / env_dir_path)
candidates.append(Path(__file__).resolve().parents[3] / env_dir_path)
candidates.append(Path(__file__).resolve().parent) candidates.append(Path(__file__).resolve().parents[3] / "target" / "release")
for directory in candidates:
for name in CDYLIB_NAMES:
p = directory / name
if p.is_file():
return p
raise LibraryNotFoundError(
"no pith-digest cdylib found (searched PITH_CDYLIB, PITH_CDYLIB_DIR, "
"the package directory and <repo>/target/release); "
"run `cargo build --release` first"
)
_lib: ctypes.CDLL | None = None
def _load() -> ctypes.CDLL:
global _lib
if _lib is None:
lib = ctypes.CDLL(str(find_cdylib()))
for name in ("pith_digest_crc32", "pith_digest_adler32"):
fn = getattr(lib, name)
fn.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_uint32)]
fn.restype = ctypes.c_int32
lib.pith_digest_fnv1a64.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.POINTER(ctypes.c_uint64),
]
lib.pith_digest_fnv1a64.restype = ctypes.c_int32
lib.pith_digest_sha256.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_char_p, ]
lib.pith_digest_sha256.restype = ctypes.c_int32
lib.pith_digest_splitmix64_fill.argtypes = [
ctypes.c_uint64,
ctypes.POINTER(ctypes.c_uint64),
ctypes.c_size_t,
]
lib.pith_digest_splitmix64_fill.restype = ctypes.c_int32
for name, out_len in (
("pith_digest_sha1", 20),
("pith_digest_sha512", 64),
):
fn = getattr(lib, name)
fn.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_char_p, ]
fn.restype = ctypes.c_int32
for name, out_len in (
("pith_digest_hmac_sha1", 20),
("pith_digest_hmac_sha256", 32),
("pith_digest_hmac_sha512", 64),
):
fn = getattr(lib, name)
fn.argtypes = [
ctypes.c_void_p, ctypes.c_size_t,
ctypes.c_void_p, ctypes.c_size_t,
ctypes.c_char_p, ]
fn.restype = ctypes.c_int32
lib.pith_digest_xxh64.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_uint64,
ctypes.POINTER(ctypes.c_uint64),
]
lib.pith_digest_xxh64.restype = ctypes.c_int32
lib.pith_digest_murmur3_x64_128.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_uint32,
ctypes.c_char_p, ]
lib.pith_digest_murmur3_x64_128.restype = ctypes.c_int32
lib.pith_digest_crc32c.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.POINTER(ctypes.c_uint32),
]
lib.pith_digest_crc32c.restype = ctypes.c_int32
for name in ("pith_digest_base64_encode", "pith_digest_base64_decode"):
fn = getattr(lib, name)
fn.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.c_char_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t), ]
fn.restype = ctypes.c_int32
lib.pith_digest_xoshiro256_fill.argtypes = [
ctypes.c_uint64,
ctypes.POINTER(ctypes.c_uint64),
ctypes.c_size_t,
]
lib.pith_digest_xoshiro256_fill.restype = ctypes.c_int32
_lib = lib
return _lib
def _bytes_op(fn: ctypes._FuncType, data: bytes, out: ctypes._CData, op: str) -> int:
status = fn(data, len(data), ctypes.byref(out))
if status != STATUS_OK:
raise FfiError(op, status)
return out.value
def sha256(data: bytes) -> bytes:
out = (ctypes.c_char * 32)()
status = _load().pith_digest_sha256(data, len(data), out)
if status != STATUS_OK:
raise FfiError("pith_digest_sha256", status)
return bytes(out)
def crc32(data: bytes) -> int:
return _bytes_op(_load().pith_digest_crc32, data, ctypes.c_uint32(), "pith_digest_crc32")
def adler32(data: bytes) -> int:
return _bytes_op(_load().pith_digest_adler32, data, ctypes.c_uint32(), "pith_digest_adler32")
def fnv1a64(data: bytes) -> int:
return _bytes_op(_load().pith_digest_fnv1a64, data, ctypes.c_uint64(), "pith_digest_fnv1a64")
def splitmix64_fill(seed: int, count: int) -> list[int]:
if count < 0:
raise ValueError("count must not be negative")
out = (ctypes.c_uint64 * count)()
status = _load().pith_digest_splitmix64_fill(seed, out, count)
if status != STATUS_OK:
raise FfiError("pith_digest_splitmix64_fill", status)
return list(out)
def _digest_op(fn: ctypes._FuncType, data: bytes, size: int, op: str) -> bytes:
out = (ctypes.c_char * size)()
status = fn(data, len(data), out)
if status != STATUS_OK:
raise FfiError(op, status)
return bytes(out)
def sha1(data: bytes) -> bytes:
return _digest_op(_load().pith_digest_sha1, data, 20, "pith_digest_sha1")
def sha512(data: bytes) -> bytes:
return _digest_op(_load().pith_digest_sha512, data, 64, "pith_digest_sha512")
def _hmac_op(fn: ctypes._FuncType, key: bytes, data: bytes, size: int, op: str) -> bytes:
out = (ctypes.c_char * size)()
status = fn(key, len(key), data, len(data), out)
if status != STATUS_OK:
raise FfiError(op, status)
return bytes(out)
def hmac_sha1(key: bytes, data: bytes) -> bytes:
return _hmac_op(_load().pith_digest_hmac_sha1, key, data, 20, "pith_digest_hmac_sha1")
def hmac_sha256(key: bytes, data: bytes) -> bytes:
return _hmac_op(_load().pith_digest_hmac_sha256, key, data, 32, "pith_digest_hmac_sha256")
def hmac_sha512(key: bytes, data: bytes) -> bytes:
return _hmac_op(_load().pith_digest_hmac_sha512, key, data, 64, "pith_digest_hmac_sha512")
def xxh64(data: bytes, seed: int = 0) -> int:
out = ctypes.c_uint64()
status = _load().pith_digest_xxh64(data, len(data), seed, ctypes.byref(out))
if status != STATUS_OK:
raise FfiError("pith_digest_xxh64", status)
return out.value
def murmur3_x64_128(data: bytes, seed: int = 0) -> bytes:
out = (ctypes.c_char * 16)()
status = _load().pith_digest_murmur3_x64_128(data, len(data), seed, out)
if status != STATUS_OK:
raise FfiError("pith_digest_murmur3_x64_128", status)
return bytes(out)
def crc32c(data: bytes) -> int:
return _bytes_op(_load().pith_digest_crc32c, data, ctypes.c_uint32(), "pith_digest_crc32c")
def _base64_op(
fn: ctypes._FuncType, data: bytes, capacity: int, op: str
) -> tuple[bytes, int]:
out = (ctypes.c_char * capacity)()
out_len = ctypes.c_size_t(0)
status = fn(data, len(data), out, capacity, ctypes.byref(out_len))
if status != STATUS_OK:
raise FfiError(op, status)
return out.raw[: out_len.value], out_len.value
def base64_encode(data: bytes) -> bytes:
capacity = 4 * ((len(data) + 2) // 3) + 1 return _base64_op(_load().pith_digest_base64_encode, data, capacity, "pith_digest_base64_encode")[0]
def base64_decode(data: bytes) -> bytes:
capacity = 3 * ((len(data) + 3) // 4) + 1
return _base64_op(_load().pith_digest_base64_decode, data, capacity, "pith_digest_base64_decode")[0]
def xoshiro256_fill(seed: int, count: int) -> list[int]:
if count < 0:
raise ValueError("count must not be negative")
out = (ctypes.c_uint64 * count)()
status = _load().pith_digest_xoshiro256_fill(seed, out, count)
if status != STATUS_OK:
raise FfiError("pith_digest_xoshiro256_fill", status)
return list(out)