File size: 9,080 Bytes
e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 e5dfae7 3bff849 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import argparse
import os
import pandas as pd
from tools.config import get_or_create_env_var, LOCAL_PII_OPTION, AWS_PII_OPTION, SELECTABLE_TEXT_EXTRACT_OPTION, TESSERACT_TEXT_EXTRACT_OPTION, TEXTRACT_TEXT_EXTRACT_OPTION
from tools.helper_functions import ensure_output_folder_exists
from tools.file_conversion import get_input_file_names, prepare_image_or_pdf
from tools.file_redaction import choose_and_run_redactor
from tools.anonymisation import anonymise_files_with_open_text
# --- Constants and Configuration ---
INPUT_FOLDER = 'input/'
OUTPUT_FOLDER = 'output/'
DEFAULT_LANGUAGE = 'en'
# Define entities for redaction
chosen_comprehend_entities = [
'BANK_ACCOUNT_NUMBER', 'BANK_ROUTING', 'CREDIT_DEBIT_NUMBER',
'CREDIT_DEBIT_CVV', 'CREDIT_DEBIT_EXPIRY', 'PIN', 'EMAIL', 'ADDRESS',
'NAME', 'PHONE', 'PASSPORT_NUMBER', 'DRIVER_ID', 'USERNAME', 'PASSWORD',
'IP_ADDRESS', 'MAC_ADDRESS', 'LICENSE_PLATE', 'VEHICLE_IDENTIFICATION_NUMBER',
'UK_NATIONAL_INSURANCE_NUMBER', 'INTERNATIONAL_BANK_ACCOUNT_NUMBER',
'SWIFT_CODE', 'UK_NATIONAL_HEALTH_SERVICE_NUMBER'
]
chosen_redact_entities = [
"TITLES", "PERSON", "PHONE_NUMBER", "EMAIL_ADDRESS", "STREETNAME", "UKPOSTCODE"
]
# --- Main CLI Function ---
def main():
"""
A unified command-line interface to prepare, redact, and anonymise various document types.
"""
parser = argparse.ArgumentParser(
description='A versatile CLI for redacting PII from PDF/image files and anonymising Word/tabular data.',
formatter_class=argparse.RawTextHelpFormatter
)
# --- General Arguments (apply to all file types) ---
general_group = parser.add_argument_group('General Options')
general_group.add_argument('--input_file', required=True, help='Path to the input file to process.')
general_group.add_argument('--output_dir', default=OUTPUT_FOLDER, help='Directory for all output files.')
general_group.add_argument('--language', default=DEFAULT_LANGUAGE, help='Language of the document content.')
general_group.add_argument('--allow_list', help='Path to a CSV file with words to exclude from redaction.')
general_group.add_argument('--pii_detector',
choices=[LOCAL_PII_OPTION, AWS_PII_OPTION],
default=LOCAL_PII_OPTION,
help='Core PII detection method (Local or AWS).')
general_group.add_argument('--aws_access_key', default='', help='Your AWS Access Key ID.')
general_group.add_argument('--aws_secret_key', default='', help='Your AWS Secret Access Key.')
# --- PDF/Image Redaction Arguments ---
pdf_group = parser.add_argument_group('PDF/Image Redaction Options (.pdf, .png, .jpg)')
pdf_group.add_argument('--ocr_method',
choices=[SELECTABLE_TEXT_EXTRACT_OPTION, TESSERACT_TEXT_EXTRACT_OPTION, TEXTRACT_TEXT_EXTRACT_OPTION],
default=TESSERACT_TEXT_EXTRACT_OPTION,
help='OCR method for text extraction from images.')
pdf_group.add_argument('--page_min', type=int, default=0, help='First page to redact.')
pdf_group.add_argument('--page_max', type=int, default=999, help='Last page to redact.')
pdf_group.add_argument('--prepare_for_review', action='store_true', help='Prepare files for reviewing redactions.')
pdf_group.add_argument('--no_images', action='store_false', dest='prepare_images', help='Disable image creation for PDF pages.')
# --- Word/Tabular Anonymisation Arguments ---
tabular_group = parser.add_argument_group('Word/Tabular Anonymisation Options (.docx, .csv, .xlsx)')
tabular_group.add_argument('--anon_strat', choices=['redact', 'encrypt', 'hash'], default='redact', help='The anonymisation strategy to apply.')
tabular_group.add_argument('--columns', nargs='+', default=[], help='A list of column names to anonymise in tabular data.')
tabular_group.add_argument('--excel_sheets', nargs='+', default=[], help='Specific Excel sheet names to process.')
tabular_group.add_argument('--deny_list', help='Path to a CSV file with specific terms/phrases to redact.')
tabular_group.add_argument('--fuzzy_mistakes', type=int, default=1, help='Number of allowed spelling mistakes for fuzzy matching.')
args = parser.parse_args()
# --- Initial Setup ---
ensure_output_folder_exists(args.output_dir)
_, file_extension = os.path.splitext(args.input_file)
file_extension = file_extension.lower()
# Load allow/deny lists
allow_list = pd.read_csv(args.allow_list) if args.allow_list else pd.DataFrame()
deny_list = pd.read_csv(args.deny_list).iloc[:, 0].tolist() if args.deny_list else []
# --- Route to the Correct Workflow Based on File Type ---
# Workflow 1: PDF/Image Redaction
if file_extension in ['.pdf', '.png', '.jpg', '.jpeg']:
print("--- Detected PDF/Image file. Starting Redaction Workflow... ---")
try:
# Step 1: Prepare the document
print("\nStep 1: Preparing document...")
(
prep_summary, prepared_pdf_paths, image_file_paths, _, _, pdf_doc,
image_annotations, _, original_cropboxes, page_sizes, textract_output_found, _, _, _, _
) = prepare_image_or_pdf(
file_paths=[args.input_file], text_extract_method=args.ocr_method,
all_line_level_ocr_results_df=pd.DataFrame(), all_page_line_level_ocr_results_with_words_df=pd.DataFrame(),
first_loop_state=True, prepare_for_review=args.prepare_for_review,
output_folder=args.output_dir, prepare_images=args.prepare_images
)
print(f"Preparation complete. {prep_summary}")
# Step 2: Redact the prepared document
print("\nStep 2: Running redaction...")
(
output_summary, output_files, _, _, log_files, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _, _
) = choose_and_run_redactor(
file_paths=[args.input_file], prepared_pdf_file_paths=prepared_pdf_paths,
pdf_image_file_paths=image_file_paths, chosen_redact_entities=chosen_redact_entities,
chosen_redact_comprehend_entities=chosen_comprehend_entities, text_extraction_method=args.ocr_method,
in_allow_list=allow_list, first_loop_state=True, page_min=args.page_min, page_max=args.page_max,
pymupdf_doc=pdf_doc, annotations_all_pages=image_annotations, page_sizes=page_sizes,
document_cropboxes=original_cropboxes, pii_identification_method=args.pii_detector,
aws_access_key_textbox=args.aws_access_key, aws_secret_key_textbox=args.aws_secret_key,
language=args.language, output_folder=args.output_dir
)
print("\n--- Redaction Process Complete ---")
print(f"Summary: {output_summary}")
print(f"\nOutput files saved to: {args.output_dir}")
print("Generated Files:", sorted(output_files))
if log_files: print("Log Files:", sorted(log_files))
except Exception as e:
print(f"\nAn error occurred during the PDF/Image redaction workflow: {e}")
# Workflow 2: Word/Tabular Data Anonymisation
elif file_extension in ['.docx', '.xlsx', '.xls', '.csv', '.parquet']:
print("--- Detected Word/Tabular file. Starting Anonymisation Workflow... ---")
try:
# Run the anonymisation function directly
output_summary, output_files, _, _, log_files, _, _ = anonymise_files_with_open_text(
file_paths=[args.input_file],
in_text="", # Not used for file-based operations
anon_strat=args.anon_strat,
chosen_cols=args.columns,
chosen_redact_entities=chosen_redact_entities,
in_allow_list=allow_list,
in_excel_sheets=args.excel_sheets,
first_loop_state=True,
output_folder=args.output_dir,
in_deny_list=deny_list,
max_fuzzy_spelling_mistakes_num=args.fuzzy_mistakes,
pii_identification_method=args.pii_detector,
chosen_redact_comprehend_entities=chosen_comprehend_entities,
aws_access_key_textbox=args.aws_access_key,
aws_secret_key_textbox=args.aws_secret_key,
language=args.language
)
print("\n--- Anonymisation Process Complete ---")
print(f"Summary: {output_summary}")
print(f"\nOutput files saved to: {args.output_dir}")
print("Generated Files:", sorted(output_files))
if log_files: print("Log Files:", sorted(log_files))
except Exception as e:
print(f"\nAn error occurred during the Word/Tabular anonymisation workflow: {e}")
else:
print(f"Error: Unsupported file type '{file_extension}'.")
print("Supported types for redaction: .pdf, .png, .jpg, .jpeg")
print("Supported types for anonymisation: .docx, .xlsx, .xls, .csv, .parquet")
if __name__ == "__main__":
main() |