99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from dictionary_matcher import match_text
|
|
from ocr_processor import analyze_screenshot
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent
|
|
DICTIONARY_DIR = BASE_DIR / "dictionaries"
|
|
REGEX_PACK_DIR = BASE_DIR / "regex-packs"
|
|
|
|
|
|
def resolve_dictionary_pack(name: str | None) -> str | None:
|
|
if not name:
|
|
return None
|
|
path = DICTIONARY_DIR / f"{name}.json"
|
|
return str(path) if path.exists() else None
|
|
|
|
|
|
def resolve_regex_pack(name: str | None) -> str | None:
|
|
if not name:
|
|
return None
|
|
path = REGEX_PACK_DIR / f"{name}.json"
|
|
return str(path) if path.exists() else None
|
|
|
|
|
|
def analyze_text_content(
|
|
text: str,
|
|
dictionary_pack: str | None = None,
|
|
regex_pack: str | None = None,
|
|
) -> dict[str, Any]:
|
|
dictionary_path = resolve_dictionary_pack(dictionary_pack)
|
|
regex_pack_path = resolve_regex_pack(regex_pack)
|
|
matches = match_text(
|
|
text=text,
|
|
dictionary_path=dictionary_path,
|
|
regex_pack_path=regex_pack_path,
|
|
)
|
|
return {
|
|
"text": text,
|
|
"dictionary_pack": dictionary_pack,
|
|
"regex_pack": regex_pack,
|
|
"dictionary_matches": matches.get("dictionary_matches", []),
|
|
"regex_matches": matches.get("regex_matches", []),
|
|
}
|
|
|
|
|
|
def analyze_artifact(
|
|
text: str | None = None,
|
|
image_path: str | None = None,
|
|
dictionary_pack: str | None = None,
|
|
regex_pack: str | None = None,
|
|
) -> dict[str, Any]:
|
|
if image_path:
|
|
dictionary_path = resolve_dictionary_pack(dictionary_pack)
|
|
regex_pack_path = resolve_regex_pack(regex_pack)
|
|
result = analyze_screenshot(
|
|
image_path=image_path,
|
|
dictionary_path=dictionary_path,
|
|
regex_pack_path=regex_pack_path,
|
|
)
|
|
result["dictionary_pack"] = dictionary_pack
|
|
result["regex_pack"] = regex_pack
|
|
result["source"] = "image"
|
|
return result
|
|
return {
|
|
**analyze_text_content(
|
|
text=text or "",
|
|
dictionary_pack=dictionary_pack,
|
|
regex_pack=regex_pack,
|
|
),
|
|
"source": "text",
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Analyze text or screenshot with DLP dictionaries/regex packs")
|
|
parser.add_argument("--text", help="Text to analyze")
|
|
parser.add_argument("--image", help="Screenshot/image path to analyze")
|
|
parser.add_argument("--dictionary-pack", default=None)
|
|
parser.add_argument("--regex-pack", default=None)
|
|
args = parser.parse_args()
|
|
|
|
result = analyze_artifact(
|
|
text=args.text,
|
|
image_path=args.image,
|
|
dictionary_pack=args.dictionary_pack,
|
|
regex_pack=args.regex_pack,
|
|
)
|
|
print(json.dumps(result, ensure_ascii=False, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|