Irannas/Masked_Email_Classification

1"""Utility functions for detecting and masking Personally Identifiable Information (PII) in text."""
2
3import re
4
5
6def mask_pii(text):
7    """
8    Detect and mask personally identifiable information (PII) in the input text.
9
10    Args:
11        text (str): The raw input text (e.g., an email body).
12
13    Returns:
14        tuple:
15            - masked_text (str): Text with PII replaced by placeholder tags.
16            - entities (list): A list of dictionaries, each containing:
17                - position (list): Start and end character positions of the PII.
18                - classification (str): Type of PII (e.g., 'email', 'phone_number').
19                - entity (str): The original PII detected.
20    """
21    entities = []
22    masked_text = text
23
24    patterns = {
25        "aadhar_num": r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",
26        "credit_debit_no": r"\b(?:\d[ -]*?){13,16}\b",
27        "expiry_no": r"\b(0[1-9]|1[0-2])\/(\d{2}|\d{4})\b",
28        "cvv_no": r"\b\d{3}\b",
29        "dob": r"\b\d{2}/\d{2}/\d{4}\b",
30        "phone_number": r"\b[6-9]\d{9}\b",
31        "email": r"\b[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}\b",
32        "full_name": (
33            r"(?i)\b(?:i\s*am|i'm|my\s*name\s*is|this\s*is)\s+"
34            r"([A-Z][a-z]+(?:\s[A-Z][a-z]+)?)"
35        ),
36    }
37
38    matches = []
39    # for label, pattern in patterns.items():
40    #     for match in re.finditer(pattern, text):
41    #         matches.append((match.start(), match.end(), label, match.group(1)))
42
43    # matches.sort(reverse=True)
44    # print(matches)
45
46    for label, pattern in patterns.items():
47        for match in re.finditer(pattern, text):
48            # If pattern includes a capturing group (like "full_name"), use the group
49            if label == "full_name" and match.lastindex:
50                start, end = match.span(1)  # span of the captured name only
51                original_value = match.group(1)
52            else:
53                start, end = match.span()  # span of the full match
54                original_value = match.group()
55            matches.append((start, end, label, original_value))
56    matches.sort(reverse=True)
57    for start, end, label, original_value in matches:
58        masked_text = masked_text[:start] + f"[{label}]" + masked_text[end:]
59        entities.append(
60            {
61                "position": [start, end],
62                "classification": label,
63                "entity": original_value,
64            }
65        )
66
67    return masked_text, entities
68