protegrity-ai-developer-python 1.2.1__py3-none-any.whl

protegrity_developer_python/utils/ccn_processing.py

@@ -0,0 +1,59 @@
+"""
+Module for processing credit card numbers (CCNs) by removing and reconstructing separators.
+"""
+
+
+def clean_ccn(ccn_string: str) -> tuple[str, dict]:
+    """
+    Remove separators from credit card number and store their positions.
+
+    Args:
+        ccn_string (str): Credit card number with separators
+
+    Returns:
+        tuple: (cleaned_string, separator_map)
+            - cleaned_string: CCN with only digits
+            - separator_map: {position: separator_char}
+    """
+    cleaned = ""
+    separator_map = {}
+
+    for i, char in enumerate(ccn_string):
+        if char.isdigit():
+            cleaned += char
+        else:
+            # Store the original position and the separator character
+            separator_map[i] = char
+
+    return cleaned, separator_map
+
+
+def reconstruct_ccn(cleaned_ccn: str, separator_map: dict) -> str:
+    """
+    Reconstruct original CCN format using the separator map.
+
+    Args:
+        cleaned_ccn (str): CCN with only digits
+        separator_map (dict): {position: separator_char} mapping
+
+    Returns:
+        str: CCN with separators restored to original positions
+    """
+    result = ""
+    original_pos = 0
+
+    for cleaned_pos, digit in enumerate(cleaned_ccn):
+        # Add any separators that should come before this digit
+        while original_pos in separator_map:
+            result += separator_map[original_pos]
+            original_pos += 1
+
+        result += digit
+        original_pos += 1
+
+    # Add any trailing separators
+    while original_pos in separator_map:
+        result += separator_map[original_pos]
+        original_pos += 1
+
+    return result

protegrity_developer_python/utils/config.py

@@ -0,0 +1,60 @@
+"""
+Module for
+"""
+
+from protegrity_developer_python.utils.constants import CONFIG as _root_config, get_config
+from typing import Dict, Optional
+from protegrity_developer_python.utils.logger import get_logger
+import logging
+
+# Get logger instance
+logger = get_logger()
+
+
+def configure(
+    endpoint_url: Optional[str] = None,
+    named_entity_map: Optional[Dict[str, str]] = None,
+    masking_char: Optional[str] = None,
+    classification_score_threshold: Optional[float] = None,
+    method: Optional[str] = None,
+    enable_logging: Optional[bool] = None,
+    log_level: Optional[str] = None,
+) -> None:
+    """
+    Configure the protegrity_developer_python module.
+
+    Args:
+        endpoint_url (str): URL of the discovery classification API.
+        named_entity_map (dict): Mapping of entity types to labels.
+        masking_char (str): Character used for masking PII.
+        classification_score_threshold (float): Minimum score to consider classification.
+        method (str): Either 'redact' or 'mask'.
+        enable_logging (bool): Enable or disable logging.
+        log_level (str): Set the logging level.
+    """
+
+    # Configure data-discovery settings
+    dd_config = get_config("data-discovery")
+    if endpoint_url:
+        dd_config["endpoint_url"] = endpoint_url
+    if named_entity_map:
+        dd_config["named_entity_map"] = named_entity_map
+    if masking_char:
+        dd_config["masking_char"] = masking_char
+    if classification_score_threshold is not None:
+        dd_config["classification_score_threshold"] = classification_score_threshold
+    if method in ("redact", "mask"):
+        dd_config["method"] = method
+    elif method:
+        logger.warning(
+            "Invalid method specified: %s. Must be 'redact' or 'mask'.", method
+        )
+    
+    # Configure global logging settings
+    if enable_logging is not None:
+        _root_config["enable_logging"] = enable_logging
+        logger.disabled = not enable_logging
+    if log_level:
+        _root_config["log_level"] = log_level
+        level = getattr(logging, log_level.upper(), logging.INFO)
+        logger.setLevel(level)

protegrity_developer_python/utils/constants.py

@@ -0,0 +1,123 @@
+"""
+Module for constants and default configurations.
+"""
+
+import os
+
+# Mapping of entity types to Protegrity data elements
+DATA_ELEMENT_MAPPING = {
+    # Account Information
+    "ACCOUNT_NAME": "string",
+    "ACCOUNT_NUMBER": "number",
+    
+    # Personal Information
+    "AGE": "number",
+    "GENDER": "string",
+    "TITLE": "string",
+    "PERSON": "string",
+    
+    # Financial
+    "AMOUNT": "number",
+    "BANK_ACCOUNT": "number",
+    "CREDIT_CARD": "ccn",
+    "CURRENCY": "string",
+    "CURRENCY_CODE": "string",
+    "CURRENCY_NAME": "string",
+    "CURRENCY_SYMBOL": "string",
+    
+    # Cryptocurrency
+    "CRYPTO_ADDRESS": "address",
+    
+    # Temporal
+    "DATETIME": "datetime",
+    "DOB": "datetime",
+    "DATE_OF_BIRTH": "datetime",
+    
+    # Identification Documents
+    "DRIVER_LICENSE": "number",
+    "PASSPORT": "passport",
+    "NATIONAL_ID": "nin",
+    "SOCIAL_SECURITY_ID": "ssn",
+    
+    # Contact Information
+    "EMAIL_ADDRESS": "email",
+    "PHONE_NUMBER": "phone",
+    
+    # Healthcare
+    "HEALTH_CARE_ID": "number",
+    
+    # Location Information
+    "LOCATION": "address",
+    "ADDRESS": "address",
+    
+    # Network
+    "IP_ADDRESS": "address",
+    "MAC_ADDRESS": "address",
+    "URL": "address",
+    
+    # Tax
+    "TAX_ID": "number",
+    
+    # India-specific
+    "IN_VEHICLE_REGISTRATION": "number",
+    "IN_VOTER": "number",
+    "IN_GSTIN": "string",
+    
+    # Security
+    "PASSWORD": "string",
+    "USERNAME": "string",
+    "USER_NAME": "string",
+    
+    # Organization
+    "ORGANIZATION": "string",
+    
+    # National Registration
+    "NRP": "number",
+    
+    # Korea
+    "KR_RRN": "number",
+    
+    # Thailand
+    "TH_TNIN": "nin",
+}
+
+CONFIG = {
+    "enable_logging": True,
+    "log_level": "INFO",
+    "data-discovery": {
+        "endpoint_url": os.getenv(
+            "DISCOVER_URL", "http://localhost:8580/pty/data-discovery/v2/classify/text"
+        ),        "transform_url": os.getenv(
+            "TRANSFORM_URL", "http://localhost:8580/pty/data-discovery/v2/transform/label"
+        ),        "named_entity_map": {},
+        "masking_char": "#",
+        "classification_score_threshold": 0.6,
+        "method": "redact",  # or "mask"
+    },
+    "semantic-guardrails": {
+        "endpoint_url": os.getenv(
+            "SEMANTIC_GUARDRAILS_URL",
+            "http://localhost:8001/pty/semantic-guardrail/v1.1",
+        ),
+    },
+}
+
+
+def get_config(section: str = None) -> dict:
+    """
+    Get configuration for a specific section or the entire config.
+    
+    Args:
+        section (str, optional): Section name ('data-discovery', 'semantic-guardrails'). 
+                                If None, returns root config.
+    
+    Returns:
+        dict: Configuration dictionary for the requested section.
+    
+    Example:
+        >>> config = get_config('data-discovery')
+        >>> method = config['method']
+    """
+    if section and section in CONFIG:
+        return CONFIG[section]
+    return CONFIG

protegrity_developer_python/utils/discover.py

@@ -0,0 +1,49 @@
+"""
+Module for discovering PII entities using a discovery API.
+"""
+
+from typing import Dict
+import json
+import requests
+from protegrity_developer_python.utils.constants import get_config
+from protegrity_developer_python.utils.logger import get_logger
+
+
+_config = get_config("data-discovery")
+# Get logger instance
+logger = get_logger()
+
+
+def discover(text: str) -> Dict:
+    """
+    Discover PII entities in the input text using the configured REST endpoint.
+
+    Args:
+        text (str): Input text to classify.
+
+    Returns:
+        dict: Full JSON response from the classification API.
+    """
+    headers = {"Content-Type": "text/plain; charset=utf-8"}
+    params = {"score_threshold": _config["classification_score_threshold"]}
+
+    try:
+        response = requests.post(
+            _config["endpoint_url"],
+            headers=headers,
+            data=text.encode("utf-8"),
+            params=params,
+            timeout=30,
+        )
+        response.raise_for_status()
+        response_json = response.json()
+        return response_json.get("classifications", {})
+    except requests.exceptions.RequestException as e:
+        logger.error("HTTP request failed: %s", e)
+        raise
+    except json.JSONDecodeError as e:
+        logger.error("Failed to decode JSON response: %s", e)
+        raise
+    except Exception as e:
+        logger.error("Unexpected error: %s", e)
+        raise

protegrity_developer_python/utils/logger.py

@@ -0,0 +1,23 @@
+"""
+Module for  logging configuration and logger retrieval.
+"""
+
+import logging
+from protegrity_developer_python.utils.constants import CONFIG as _config
+
+logger = logging.getLogger("protegrity_developer_python")
+handler = logging.StreamHandler()
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+handler.setFormatter(formatter)
+logger.addHandler(handler)
+
+if _config["enable_logging"]:
+    level = getattr(logging, _config["log_level"].upper(), logging.INFO)
+    logger.setLevel(level)
+    logger.propagate = False
+else:
+    logger.disabled = True
+
+
+def get_logger():
+    return logger

protegrity_developer_python/utils/pii_processing.py

@@ -0,0 +1,291 @@
+"""
+Module for processing PII entities in text, including redaction, masking, and protection.
+"""
+
+import re
+from typing import Dict, Tuple
+from protegrity_developer_python.utils.ccn_processing import clean_ccn, reconstruct_ccn
+from protegrity_developer_python.utils.protector import get_protector_session
+from protegrity_developer_python.utils.logger import get_logger
+from protegrity_developer_python.utils.constants import (
+    DATA_ELEMENT_MAPPING as entity_endpoint_mapped,
+    get_config,
+)
+
+# Get logger instance
+logger = get_logger()
+
+# Get data-discovery sub-config for backward compatibility
+_config = get_config("data-discovery")
+
+
+def _merge_overlapping_entities(
+    entity_spans: Dict[Tuple[int, int], Tuple[str, float]],
+) -> Dict[Tuple[int, int], Tuple[str, float]]:
+    """
+    Merge overlapping entity spans in a dictionary.
+
+    Given a dictionary with (start, end) index tuples as keys and (entity, score) tuples as values,
+    this function merges overlapping or adjacent spans by combining entity labels with '|' and
+    retaining the highest score.
+
+    Args:
+        entity_dict (dict): Mapping of index spans to (entity, score).
+
+    Returns:
+        dict: Merged spans with combined entities and max scores.
+
+    Example:
+        Input:
+            {
+                (10, 20): ("PERSON", 0.9),
+                (15, 25): ("LOCATION", 0.85)
+            }
+        Output:
+            {
+                (10, 25): ("PERSON|LOCATION", 0.9)
+            }
+    """
+    merged = []
+    for (start, end), (entity, score) in entity_spans.items():
+        if not merged:
+            merged.append([(start, end), entity, score])
+        else:
+            last_start, last_end = merged[-1][0]
+            # Check for overlap
+            if start <= last_end and end >= last_start:
+                # Merge ranges
+                new_start = min(start, last_start)
+                new_end = max(end, last_end)
+                # Merge entities (sorted by score descending, then lexicographically)
+                if entity != merged[-1][1]:
+                    last_entity = merged[-1][1]
+                    last_score = merged[-1][2]
+                    # Sort by score (descending), then lexicographically
+                    if last_score > score:
+                        entities = [last_entity, entity]
+                    elif last_score < score:
+                        entities = [entity, last_entity]
+                    else:
+                        # Scores equal, sort lexicographically
+                        entities = sorted([last_entity, entity])
+                    new_entity = "|".join(entities)
+                else:
+                    new_entity = entity
+                # Take max score
+                new_score = max(score, merged[-1][2])
+                merged[-1] = [(new_start, new_end), new_entity, new_score]
+            else:
+                merged.append([(start, end), entity, score])
+
+    # Convert back to dictionary
+    merged_entities = {tuple(k): (v, s) for k, v, s in merged}
+    logger.debug("Before Merging Entity spans: \n%s", entity_spans)
+    logger.debug("Merged Entity spans: \n%s", merged_entities)
+    return merged_entities
+
+
+def collect_entity_spans(entities: Dict) -> Dict[Tuple[int, int], Tuple[str, float]]:
+    """
+    Collects entity spans for redaction or masking based on score.
+
+    Returns a dictionary of (start_index, end_index) : (entity_name, score).
+    The dictionary is sorted in reverse order of start_index to ensure that replacements
+    do not affect the character positions of subsequent spans.
+
+    Args:
+        entities (dict): Dictionary of detected PII entities.
+
+    Returns:
+        dict: Sorted entity spans mapped to entity names and scores.
+    """
+    entity_map = {}
+    for entity_name, entity_details in entities.items():
+        for obj in entity_details:
+            score = obj.get("score")
+            loc = obj.get("location", {})
+            start = loc.get("start_index", 0)
+            end = loc.get("end_index", 0)
+            if (start, end) in entity_map:
+                if score > entity_map[(start, end)][1]:
+                    entity_map[(start, end)] = (entity_name, score)
+            else:
+                entity_map[(start, end)] = (entity_name, score)
+    # Sort entity_map in reverse order of start index to avoid index shifting
+    entity_spans = {
+        key: entity_map[key]
+        for key in sorted(entity_map, key=lambda x: x[0], reverse=True)
+    }
+
+    logger.debug("Raw Entity spans collected: \n%s", entity_spans)
+    return _merge_overlapping_entities(entity_spans)
+
+def unprotect_data(text: str) -> str:
+    session = get_protector_session()
+    try:
+        # Regex pattern to match [entity]data[/entity]
+        pattern = r'\[([^\]]+)\]([^[]*)\[\/\1\]'
+        
+        def replace_protected_data(match):
+            data_element = entity_endpoint_mapped[match.group(1)] 
+            entity = match.group(1)
+            data = match.group(2)
+            try:
+                unprotect_dataprotected = ""
+                if entity == "CREDIT_CARD":
+                    input_data, separator_map = clean_ccn(data)
+                    unprotect_dataprotected = session.unprotect(input_data, data_element)
+                    unprotect_dataprotected = reconstruct_ccn(unprotect_dataprotected, separator_map)
+                else:
+                    unprotect_dataprotected = session.unprotect(data, data_element)
+                logger.info(f"Unprotected {data_element}: {data}")
+                return unprotect_dataprotected  # Return just the data, removing the tags
+            except Exception:
+                logger.warning(f"Failed to unprotect {data_element}: {data}.")
+                return data  # Return original data if unprotection fails
+        
+        # Replace all matches with unprotected data
+        result = re.sub(pattern, replace_protected_data, text)
+        return result
+    except Exception as e:
+        logger.error("Failed to process text: %s", e)
+        raise
+
+def protect_data(
+    merged_entities: Dict[Tuple[int, int], Tuple[str, float]], text: str
+) -> str:
+    """
+    Protects the identified PII entities in the input text.
+
+    Args:
+        entities (dict): Dictionary of detected PII entities.
+        text (str): Original input text.
+
+    Returns:
+        text (str): Protected PII text.
+    """
+    session = get_protector_session()
+    for key, val in merged_entities.items():
+        start, end = key
+        entity, _ = val
+        data_element = None
+        protected = ""
+        get_named_entity_map = None
+        entity_selected = None
+
+        if "|" in entity:
+            data_element = entity_endpoint_mapped[entity.split("|")[0]]
+            entity_selected = entity.split("|")[0]
+            logger.debug("Multiple entities '%s' detected at span [%d:%d], using first entity for protection - %s.", entity, start, end, entity_selected)
+            get_named_entity_map = _config["named_entity_map"].get(entity.split("|")[0], None)
+        else:
+            data_element = entity_endpoint_mapped[entity]
+            entity_selected = entity
+            get_named_entity_map = _config["named_entity_map"].get(entity, None)
+        input_data = text[start:end]
+
+        if get_named_entity_map:
+            try:
+                if entity == "CREDIT_CARD":
+                    input_data, separator_map = clean_ccn(input_data)
+                    protected = session.protect(input_data, data_element)
+                    protected = reconstruct_ccn(protected, separator_map)
+                else:
+                    protected = session.protect(input_data, data_element)
+
+                logger.info(
+                    "Entity '%s' at span [%d:%d] protected successfully.",
+                    entity,
+                    start,
+                    end,
+                )
+
+                text = text[:start] + f"[{entity_selected}]" + protected + f"[/{entity_selected}]" + text[end:]
+                logger.debug(
+                    "Entity '%s' at span [%d:%d] protected as [%s]%s[/%s].",
+                    entity,
+                    start,
+                    end,
+                    entity_selected,
+                    protected,
+                    entity_selected,
+                )
+            except Exception:
+                logger.warning(
+                    "Failed to protect entity '%s' at span [%d:%d] having value %s.",
+                    entity,
+                    start,
+                    end,
+                    input_data,
+                )
+                text = text[:start] + input_data + text[end:]
+        else:
+            logger.warning(
+                "Entity '%s' detected at span [%d:%d] but no mapping found in named_entity_map - skipping protection",
+                entity,
+                start,
+                end,
+            )
+
+    return text
+
+
+def redact_data(
+    merged_entities: Dict[Tuple[int, int], Tuple[str, float]], text: str
+) -> str:
+    """
+    Redacts the identified PII entities in the input text.
+
+    Args:
+        entities (dict): Dictionary of detected PII entities.
+        text (str): Original input text.
+
+    Returns:
+        text (str): Redacted PII text.
+    """
+    for key, val in merged_entities.items():
+        start, end = key
+        entity_name = val[0]
+        logger.debug(
+            "Entity '%s' found at span [%d:%d] with score %f",
+            entity_name,
+            start,
+            end,
+            val[1],
+        )
+        if _config["method"] == "redact":
+            if "|" in entity_name:
+                label = "|".join(
+                    _config["named_entity_map"].get(entity, entity)
+                    for entity in entity_name.split()
+                )
+            else:
+                label = _config["named_entity_map"].get(entity_name, "")
+            if label:
+                label = f"[{label}]"
+                logger.info(
+                    "Entity '%s' found at span [%d:%d]... redacted as %s",
+                    entity_name,
+                    start,
+                    end,
+                    label, 
+                )
+                text = text[:start] + label + text[end:]
+            else:
+                label = f"[{entity_name}]"
+                logger.warning(
+                    "Entity '%s' detected at span [%d:%d] but no mapping found in named_entity_map - skipping redaction",
+                    entity_name,
+                    start,
+                    end,
+                )
+        elif _config["method"] == "mask":
+            logger.info(
+                "Entity '%s' found at span [%d:%d]... masked",
+                entity_name,
+                start,
+                end,
+            )
+            text = text[:start] + _config["masking_char"] * (end - start) + text[end:]
+
+    return text

protegrity_developer_python/utils/protector.py

@@ -0,0 +1,23 @@
+"""
+Module for interfacing with Protegrity's appython Protector for data protection.
+"""
+
+from appython import Protector
+
+instance, session = None, None
+
+
+# Initialize the Protector instance
+def _initialize_protector():
+    global instance
+    instance = Protector()
+    return instance
+
+
+# Get or create a Protector session
+def get_protector_session():
+    global instance, session
+    if instance is None:
+        protector = _initialize_protector()
+        session = protector.create_session("superuser")
+    return session