105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Filesystem and I/O helpers for the intrusion-detection system."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
import gzip
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Any, Iterable, Iterator, Sequence
|
|
|
|
ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")
|
|
|
|
|
|
# -- directory --------------------------------------------------------------
|
|
|
|
def validate_dir(path: str, need_write: bool = False) -> str:
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(f"path does not exist: {path}")
|
|
if not os.path.isdir(path):
|
|
raise NotADirectoryError(f"path is not a directory: {path}")
|
|
if not os.access(path, os.R_OK):
|
|
raise PermissionError(f"path not readable: {path}")
|
|
if need_write and not os.access(path, os.W_OK):
|
|
raise PermissionError(f"path not writable: {path}")
|
|
return path
|
|
|
|
|
|
def ensure_dirs(*paths: str) -> None:
|
|
for p in paths:
|
|
os.makedirs(p, exist_ok=True)
|
|
|
|
|
|
# -- log discovery ----------------------------------------------------------
|
|
|
|
def find_logs(path: str, pattern: re.Pattern | None = None,
|
|
suffixes: Sequence[str] = (),
|
|
exclude_prefixes: Sequence[str] = ()) -> list[str]:
|
|
"""List matching log file names (not full paths), sorted."""
|
|
out = []
|
|
for name in os.listdir(path):
|
|
if not os.path.isfile(os.path.join(path, name)):
|
|
continue
|
|
if any(name.startswith(p) for p in exclude_prefixes):
|
|
continue
|
|
if pattern is not None and not pattern.match(name):
|
|
continue
|
|
if suffixes and not name.endswith(tuple(suffixes)):
|
|
continue
|
|
out.append(name)
|
|
return sorted(out)
|
|
|
|
|
|
def read_lines(path: str) -> Iterator[str]:
|
|
"""Yield ANSI-stripped, newline-stripped lines. Transparently gunzips."""
|
|
opener = gzip.open if path.endswith(".gz") else open
|
|
with opener(path, "rt", encoding="utf-8", errors="replace") as f:
|
|
for line in f:
|
|
yield ANSI_RE.sub("", line.rstrip("\r\n"))
|
|
|
|
|
|
def read_lines_from(path: str, start: int = 0) -> Iterator[tuple[str, int]]:
|
|
"""Yield ``(line, resume_offset)`` for a plain-text log from byte ``start``.
|
|
|
|
Only seekable (non-gzip) files can be resumed. ``resume_offset`` points
|
|
after the last complete line; a trailing partial line (no newline yet) is
|
|
yielded but resets the resume offset to its start so it is re-read once
|
|
completed. The last yielded ``resume_offset`` is the value to store.
|
|
"""
|
|
with open(path, "rb") as f:
|
|
f.seek(start)
|
|
pos = start
|
|
for raw in f:
|
|
line_start = pos
|
|
pos += len(raw)
|
|
complete = raw.endswith(b"\n") or raw.endswith(b"\r")
|
|
text = ANSI_RE.sub(
|
|
"", raw.decode("utf-8", errors="replace").rstrip("\r\n"))
|
|
yield text, (pos if complete else line_start)
|
|
|
|
|
|
# -- writers ----------------------------------------------------------------
|
|
|
|
def write_csv(path: str, header: Sequence[str],
|
|
rows: Iterable[Sequence[Any]]) -> str:
|
|
with open(path, "w", newline="", encoding="utf-8") as f:
|
|
w = csv.writer(f)
|
|
w.writerow(header)
|
|
w.writerows(rows)
|
|
return path
|
|
|
|
|
|
def write_json(path: str, obj: Any) -> str:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(obj, f, indent=2, default=str)
|
|
return path
|
|
|
|
|
|
def write_lines(path: str, lines: Iterable[str]) -> str:
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
for line in lines:
|
|
f.write(line + "\n")
|
|
return path
|