dataRetreiver.py

#!/usr/bin/env python3
"""
Data Retriever Script for MySQL Database
Connects to MySQL, retrieves account data, parses DG1 and DG2, and outputs to CSV
"""

import argparse
import csv
import os
import sys
from binascii import unhexlify
from datetime import datetime
import mysql.connector
from mysql.connector import Error

# Import DG1 parsing logic
from binascii import unhexlify
try:
    from mrz.checker.td3 import TD3CodeChecker as TD3Code
except ImportError:
    try:
        from mrz.checker.td3 import TD3Code
    except ImportError:
        try:
            from mrz import TD3
            TD3Code = TD3
        except ImportError:
            print("Error: Could not import MRZ library. Please install it with: pip3 install mrz")
            sys.exit(1)


def parse_arguments():
    """Parse command line arguments for MySQL connection parameters"""
    parser = argparse.ArgumentParser(description='Retrieve and parse account data from MySQL')
    parser.add_argument('--host', required=True, help='MySQL server host')
    parser.add_argument('--port', type=int, default=3306, help='MySQL server port (default: 3306)')
    parser.add_argument('--username', required=True, help='MySQL username')
    parser.add_argument('--password', required=True, help='MySQL password')
    parser.add_argument('--database', required=True, help='MySQL database name')
    parser.add_argument('--output', default='account_data.csv', help='Output CSV file (default: account_data.csv)')
    parser.add_argument('--images-dir', default='personalImages', help='Directory for extracted images (default: personalImages)')
    parser.add_argument('--uid-filter', help='Comma-separated list of UIDs to filter. Use quotes for UIDs with special characters (e.g., "1,2,3" or \'user@domain.com,admin.test\'). Supports both numeric and string UIDs. If not provided, all accounts will be retrieved.')
    parser.add_argument('--expired-before', help='Filter users with passports expired before this date (YYYY-MM-DD format). Users will be saved in userWithExpiredPassport subfolder.')
    
    return parser.parse_args()


def connect_to_mysql(host, port, username, password, database):
    """Connect to MySQL database and return connection object"""
    try:
        print(f"Connecting to MySQL database: {database} on {host}:{port} as {username}")
        connection = mysql.connector.connect(
            host=host,
            port=port,
            user=username,
            password=password,
            database=database
        )
        if connection.is_connected():
            print(f"Successfully connected to MySQL database: {database}")
            return connection
    except Error as e:
        print(f"Error connecting to MySQL: {e}")
        sys.exit(1)


def fetch_account_data(connection, uid_filter=None):
    """Fetch data from account table, optionally filtered by UIDs"""
    try:
        cursor = connection.cursor()
        
        if uid_filter:
            # Parse comma-separated UIDs and create placeholders for SQL IN clause
            # Handle quoted UIDs properly by stripping outer quotes if present
            uid_list = []
            for uid in uid_filter.split(','):
                uid = uid.strip()
                # Remove outer quotes if present (both single and double quotes)
                if (uid.startswith('"') and uid.endswith('"')) or (uid.startswith("'") and uid.endswith("'")):
                    uid = uid[1:-1]
                if uid:  # Only add non-empty UIDs
                    uid_list.append(uid)
            
            if not uid_list:
                print("Warning: No valid UIDs provided in filter, retrieving all records")
                uid_filter = None
        
        if uid_filter and uid_list:
            # Create placeholders for SQL IN clause
            placeholders = ','.join(['%s'] * len(uid_list))
            query = f"""
            SELECT uid, country, sodId, expires, aaPublicKey, aaSigAlgo, aaCount, 
                   aaLastAuthn, dg1, dg2 
            FROM account
            WHERE uid IN ({placeholders})
            """
            cursor.execute(query, uid_list)
            print(f"Filtering by UIDs: {', '.join(uid_list)}")
        else:
            query = """
            SELECT uid, country, sodId, expires, aaPublicKey, aaSigAlgo, aaCount, 
                   aaLastAuthn, dg1, dg2 
            FROM account
            """
            cursor.execute(query)
            print("Retrieving all records from account table")
        
        records = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        cursor.close()
        
        print(f"Retrieved {len(records)} records from account table")
        return records, columns
        
    except Error as e:
        print(f"Error fetching data: {e}")
        return [], []


def mrz_date_to_readable(mrz_date: str) -> str:
    """Convert MRZ date format to readable date"""
    if not mrz_date or mrz_date == "Invalid":
        return "Invalid"
    try:
        yy = int(mrz_date[:2])
        mm = int(mrz_date[2:4])
        dd = int(mrz_date[4:6])
        current_year = datetime.now().year % 100
        
        # For birth dates: if year > current year, assume 1900s, else 2000s
        # For expiry dates: if year < current year, assume 2000s, else 1900s (but this is rare)
        # Simple heuristic: years 00-50 are 2000s, 51-99 are 1900s
        if yy <= 50:
            century = 2000
        else:
            century = 1900
            
        year = century + yy
        return datetime(year, mm, dd).strftime("%Y-%m-%d")
    except (ValueError, IndexError):
        return "Invalid"


def parse_dg1(dg1_hex):
    """Parse DG1 data to extract MRZ information"""
    if not dg1_hex:
        return {}
    
    try:
        # Handle different data types
        if isinstance(dg1_hex, bytes):
            # Already bytes, use directly
            dg1_bytes = dg1_hex
        elif isinstance(dg1_hex, str):
            # Clean and validate hex string
            dg1_hex_clean = dg1_hex.strip()
            
            # Remove any non-hex characters
            dg1_hex_clean = ''.join(c for c in dg1_hex_clean if c in '0123456789ABCDEFabcdef')
            
            # Ensure even length (pad with 0 if odd)
            if len(dg1_hex_clean) % 2 != 0:
                dg1_hex_clean = '0' + dg1_hex_clean
                
            if not dg1_hex_clean:
                print("No valid hex data found in DG1")
                return {}
            
            # Convert hex to bytes
            dg1_bytes = unhexlify(dg1_hex_clean)
        else:
            print(f"Unsupported DG1 data type: {type(dg1_hex)}")
            return {}
        
        # Extract ASCII MRZ from DG1 bytes (skip initial TLV/ASN.1 bytes)
        mrz_bytes = dg1_bytes[5:]
        
        # Filter ASCII characters (32-126)
        mrz_text = "".join(chr(b) for b in mrz_bytes if 32 <= b <= 126)
        
        # Split into lines (TD3: 2 lines × 44 chars)
        line_length = 44
        mrz_lines = [mrz_text[i:i+line_length] for i in range(0, len(mrz_text), line_length)]
        mrz_text_clean = "\n".join(mrz_lines[:2])  # Take first 2 lines
        
        if not mrz_text_clean.strip():
            return {}
        
        # Parse MRZ using mrz library
        mrz = TD3Code(mrz_text_clean)
        
        # Initialize with default values
        parsed_data = {
            "document_type": 'N/A',
            "issuing_state": 'N/A',
            "surname": 'N/A',
            "given_names": 'N/A',
            "document_number": 'N/A',
            "nationality": 'N/A',
            "date_of_birth": 'Invalid',
            "sex": 'N/A',
            "date_of_expiry": 'Invalid'
        }
        
        # Try fields attribute/method (most reliable for this MRZ library)
        if hasattr(mrz, 'fields'):
            try:
                if callable(mrz.fields):
                    fields_data = mrz.fields()
                else:
                    fields_data = mrz.fields
                
                parsed_data.update({
                    "document_type": getattr(fields_data, 'document_type', 'N/A'),
                    "issuing_state": getattr(fields_data, 'country', 'N/A'),
                    "surname": getattr(fields_data, 'surname', 'N/A'),
                    "given_names": getattr(fields_data, 'name', getattr(fields_data, 'given_names', 'N/A')),
                    "document_number": getattr(fields_data, 'document_number', 'N/A'),
                    "nationality": getattr(fields_data, 'nationality', 'N/A'),
                    "date_of_birth": mrz_date_to_readable(str(getattr(fields_data, 'birth_date', ''))),
                    "sex": getattr(fields_data, 'sex', 'N/A'),
                    "date_of_expiry": mrz_date_to_readable(str(getattr(fields_data, 'expiry_date', '')))
                })
            except Exception:
                pass
        
        return parsed_data
        
    except Exception as e:
        print(f"Error parsing DG1: {e}")
        return {}


def parse_dg2_and_save_image(dg2_hex, document_number, images_dir):
    """Parse DG2 data to extract and save image"""
    if not dg2_hex or not document_number:
        return ""
    
    try:
        # Handle different data types
        if isinstance(dg2_hex, bytes):
            # Already bytes, use directly
            dg2_bytes = dg2_hex
        elif isinstance(dg2_hex, str):
            # Clean and validate hex string
            dg2_hex_clean = dg2_hex.strip()
            
            # Remove any non-hex characters
            dg2_hex_clean = ''.join(c for c in dg2_hex_clean if c in '0123456789ABCDEFabcdef')
            
            # Ensure even length (pad with 0 if odd)
            if len(dg2_hex_clean) % 2 != 0:
                dg2_hex_clean = '0' + dg2_hex_clean
                
            if not dg2_hex_clean:
                print(f"No valid hex data found in DG2 for document {document_number}")
                return ""
            
            # Convert hex to bytes
            dg2_bytes = unhexlify(dg2_hex_clean)
        else:
            print(f"Unsupported DG2 data type: {type(dg2_hex)}")
            return ""
        
        # Create directory for this document
        doc_dir = os.path.join(images_dir, str(document_number))
        os.makedirs(doc_dir, exist_ok=True)
        
        image_path = ""
        
        # Search for JPEG image
        if b'\xff\xd8' in dg2_bytes:  # JPEG start marker
            start = dg2_bytes.find(b'\xff\xd8')
            end = dg2_bytes.find(b'\xff\xd9')
            if end != -1:
                end += 2  # Include JPEG end marker
                jpeg_bytes = dg2_bytes[start:end]
                
                image_path = os.path.join(doc_dir, f"{document_number}.jpg")
                with open(image_path, "wb") as f:
                    f.write(jpeg_bytes)
                print(f"Saved JPEG image: {image_path}")
        
        # Search for JPEG2000 image
        elif b'\x00\x00\x00\x0c' in dg2_bytes and b'jP  ' in dg2_bytes:  # JPEG2000 magic
            start = dg2_bytes.find(b'jP  ') - 4
            if start >= 0:
                jpeg2000_bytes = dg2_bytes[start:]
                
                image_path = os.path.join(doc_dir, f"{document_number}.jp2")
                with open(image_path, "wb") as f:
                    f.write(jpeg2000_bytes)
                print(f"Saved JPEG2000 image: {image_path}")
        
        return image_path
        
    except Exception as e:
        print(f"Error parsing DG2 for document {document_number}: {e}")
        return ""


def is_passport_expired_before(expiry_date_str, cutoff_date_str):
    """Check if passport expiry date is before the given cutoff date"""
    if not expiry_date_str or expiry_date_str == "Invalid" or not cutoff_date_str:
        return False
    
    try:
        # Parse expiry date (format: YYYY-MM-DD)
        expiry_date = datetime.strptime(expiry_date_str, "%Y-%m-%d")
        cutoff_date = datetime.strptime(cutoff_date_str, "%Y-%m-%d")
        return expiry_date < cutoff_date
    except ValueError:
        return False


def process_records_to_csv(records, columns, output_file, images_dir, expired_before=None, host=None, port=None, username=None, database=None):
    
    print (f"Processing records to CSV: {output_file}")
    print (f"Images directory: {images_dir}")
    print (f"Expired before: {expired_before}")
    
    # Create images directory
    os.makedirs(images_dir, exist_ok=True)
    
    # If filtering by expired passports, create subfolder
    if expired_before:
        expired_images_dir = os.path.join(images_dir, 'userWithExpiredPassport')
        os.makedirs(expired_images_dir, exist_ok=True)
        print(f'Created expired passport subfolder: {expired_images_dir}')
    
    # Prepare CSV headers
    csv_headers = [
        'uid', 'country', 'expires', 'aaSigAlgo', 
        'aaCount', 'aaLastAuthn', 'image_path',
        # DG1 parsed fields
        'document_type', 'issuing_state', 'surname', 'given_names', 
        'document_number', 'nationality', 'date_of_birth', 'sex', 'date_of_expiry'
    ]
    
    processed_data = []
    expired_users = []
    
    for record in records:
        # Convert record to dictionary
        row_data = dict(zip(columns, record))
        
        # Parse DG1
        dg1_data = parse_dg1(row_data.get('dg1', ''))
        
        # Check if this is an expired passport user
        is_expired = False
        if expired_before and dg1_data.get('date_of_expiry') and dg1_data.get('date_of_expiry') != 'Invalid':
            is_expired = is_passport_expired_before(dg1_data.get('date_of_expiry'), expired_before)
            
            if is_expired:
                # Add to expired users list for console output
                full_name = f"{dg1_data.get('surname', 'N/A')}, {dg1_data.get('given_names', 'N/A')}"
                expired_users.append({
                    'uid': row_data.get('uid', 'unknown'),
                    'name': full_name,
                    'expiry_date': dg1_data.get('date_of_expiry', 'N/A')
                })
        
        # Parse DG2 and save image
        document_number = dg1_data.get('document_number', 'N/A')
        if document_number == 'N/A':
            # Use UID as fallback for folder name
            document_number = str(row_data.get('uid', 'unknown'))
        
        # Choose the appropriate images directory
        current_images_dir = expired_images_dir if (expired_before and is_expired) else images_dir
        
        image_path = parse_dg2_and_save_image(row_data.get('dg2', ''), document_number, current_images_dir)
        
        # Update image path to reflect the actual location
        if expired_before and is_expired:
            # Update image path to include the subfolder
            if image_path:
                image_path = image_path.replace(images_dir, os.path.join(images_dir, 'userWithExpiredPassport'))
        
        print(f"Processed record UID: {row_data.get('uid', 'unknown')}, Document: {document_number}")
        
        # Prepare CSV row
        csv_row = {
            'uid': row_data.get('uid', ''),
            'country': row_data.get('country', ''),
            'expires': row_data.get('expires', ''),
            'aaSigAlgo': row_data.get('aaSigAlgo', ''),
            'aaCount': row_data.get('aaCount', ''),
            'aaLastAuthn': row_data.get('aaLastAuthn', ''),
            'image_path': image_path,
            # DG1 parsed data
            'document_type': dg1_data.get('document_type', ''),
            'issuing_state': dg1_data.get('issuing_state', ''),
            'surname': dg1_data.get('surname', ''),
            'given_names': dg1_data.get('given_names', ''),
            'document_number': dg1_data.get('document_number', ''),
            'nationality': dg1_data.get('nationality', ''),
            'date_of_birth': dg1_data.get('date_of_birth', ''),
            'sex': dg1_data.get('sex', ''),
            'date_of_expiry': dg1_data.get('date_of_expiry', '')
        }
        
        processed_data.append(csv_row)
    
    # Print expired users to console
    if expired_before and expired_users:
        print('=' * 60)
        print(f'\n\n\n\n=== USERS WITH EXPIRED PASSPORTS (expired before {expired_before}) ===')
        print('=' * 60)
        for user in expired_users:
            print(f"UID: {user['uid']} | Name: {user['name']} | Expiry Date: {user['expiry_date']}")
        print(f'Total expired users found: {len(expired_users)}')
        print('=' * 60)
        print('\n')

        print (f"To delete the expired passports from the database, run the following command: \n")
        print(f"python3 dataRemover.py --host <host> --port <port> --username <username> --password <password> --database <database> --uid-filter <list of uids to delete, separated by commas>\n")
        print(f"OR with real values: \n")
        # Generate the actual UID list from expired_users
        expired_uid_list = ','.join([str(user['uid']) for user in expired_users])
        print(f"python3 dataRemover.py --host {host} --port {port} --username {username} --password <write your password here> --database {database} --uid-filter '{expired_uid_list}'\n")

        #TODO: add the command to delete the expired passports from the database
        print('=' * 60)
    elif expired_before:
        print(f'\nNO USERS FOUND WITH PASSPORTS EXPIRED BEFORE {expired_before}')
    
    # Write to CSV
    with open(output_file, 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=csv_headers)
        writer.writeheader()
        writer.writerows(processed_data)
    
    print(f"Data written to CSV file: {output_file}")
    print(f"Total records processed: {len(processed_data)}")


def main():
    """Main function"""
    # Parse command line arguments
    args = parse_arguments()
    
    # Validate parameters
    if not all([args.host, args.username, args.password, args.database]):
        print("Error: Missing required MySQL connection parameters")
        sys.exit(1)
    
    # Connect to MySQL
    connection = connect_to_mysql(
        args.host, args.port, args.username, 
        args.password, args.database
    )
    
    try:
        # Fetch data from account table
        records, columns = fetch_account_data(connection, args.uid_filter)
        
        if not records:
            print("NO RECORDS FOUND IN ACCOUNT TABLE")
            return
        
        # Process records and generate CSV
        process_records_to_csv(records, columns, args.output, args.images_dir, args.expired_before, args.host, args.port, args.username, args.database)
        
    finally:
        # Close connection
        if connection.is_connected():
            connection.close()
            print("MySQL connection closed")


if __name__ == "__main__":
    main()