testWriteData.py

#!/usr/bin/env python3
"""
testWriteData.py - Test data writing module for Data Wiping Toolkit
Can be used both as standalone script and as imported module
"""
import sys
import subprocess
import json
import time
from pathlib import Path
import os
import hashlib
import tempfile
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')
logger = logging.getLogger(__name__)

class TestWriteError(Exception):
    """Custom exception for test write operations"""
    pass

def is_mount_point(path):
    """Check if path is a mount point"""
    try:
        return os.path.ismount(path)
    except (OSError, TypeError):
        return False

def sha256_file(path, read_bytes=None):
    """
    Compute SHA256 checksum of a file or raw device (optionally only part).
    Args:
        path (str): File path or device path
        read_bytes (int): Number of bytes to read (None for entire file)
    Returns:
        str: SHA256 hexdigest
    """
    try:
        h = hashlib.sha256()
        with open(path, "rb") as f:
            if read_bytes:
                data = f.read(read_bytes)
                if len(data) != read_bytes:
                    logger.warning(f"Only read {len(data)} bytes instead of {read_bytes}")
                h.update(data)
            else:
                for chunk in iter(lambda: f.read(4096), b""):
                    h.update(chunk)
        return h.hexdigest()
    except (OSError, IOError) as e:
        raise TestWriteError(f"Error reading {path} for checksum: {e}")

def write_safe_file(mount_point, pattern, output_dir=None):
    """
    Write test data safely to a mounted filesystem
    Args:
        mount_point (str): Path to mounted filesystem
        pattern (str): Data pattern to write
        output_dir (str): Directory to store temporary files
    Returns:
        tuple: (test_file_path, bytes_written)
    """
    try:
        test_file_path = os.path.join(mount_point, "secret_test_data.txt")
        content = (pattern + "\n") * 100
        
        with open(test_file_path, "w") as f:
            f.write(content)
        
        bytes_written = len(content.encode('utf-8'))
        logger.info(f"Test data written safely to {test_file_path} ({bytes_written} bytes)")
        return test_file_path, bytes_written
        
    except (OSError, IOError) as e:
        raise TestWriteError(f"Error writing to mount point {mount_point}: {e}")

def write_raw_dd(device, pattern, output_dir=None):
    """
    Write test data to raw device using dd
    Args:
        device (str): Device path like /dev/sda
        pattern (str): Data pattern to write
        output_dir (str): Directory to store temporary files
    Returns:
        tuple: (temp_file_path, bytes_written)
    """
    try:
        # Create temporary file in specified directory or system temp
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            temp_file_path = os.path.join(output_dir, "secret_test_data.txt")
        else:
            temp_file_path = "/tmp/secret_test_data.txt"
        
        content = (pattern + "\n") * 100
        
        with open(temp_file_path, "w") as f:
            f.write(content)
        
        bytes_written = 1024 * 1024  # 1MB
        
        logger.info(f"Writing test data to raw device {device} (1MB at start)...")
        
        # Use dd to write to raw device
        cmd = ["dd", f"if={temp_file_path}", f"of={device}", 
               "bs=1M", "count=1", "conv=fsync"]
        
        result = subprocess.run(cmd, check=True, capture_output=True, text=True)
        logger.debug(f"dd output: {result.stderr}")
        logger.info("Raw test data written successfully!")
        
        return temp_file_path, bytes_written
        
    except subprocess.CalledProcessError as e:
        raise TestWriteError(f"Error writing to raw device {device}: {e}")
    except (OSError, IOError) as e:
        raise TestWriteError(f"Error creating temporary file: {e}")

def test_write_data(target, pattern=None, output_dir=None, metadata_file=None):
    """
    Main function to write test data and verify
    Args:
        target (str): Device path or mount point
        pattern (str): Data pattern to write
        output_dir (str): Directory for temporary files
        metadata_file (str): Path to save metadata JSON
    Returns:
        dict: Test results metadata
    """
    if pattern is None:
        pattern = "SUPER_SECRET_DATA_123456789"
    
    if metadata_file is None:
        metadata_file = os.path.join(output_dir or ".", "test_write_metadata.json")
    
    try:
        start_time = time.time()
        
        # Determine write method and execute
        if is_mount_point(target):
            method = "safe filesystem write"
            file_path, bytes_written = write_safe_file(target, pattern, output_dir)
            checksum_target = file_path
            checksum_bytes = None
        else:
            method = "dd 1MB write at start of raw device"
            file_path, bytes_written = write_raw_dd(target, pattern, output_dir)
            checksum_target = target
            checksum_bytes = 1024 * 1024
        
        # Calculate checksums
        logger.info("Calculating pre-write checksum...")
        pre_checksum = sha256_file(checksum_target, read_bytes=checksum_bytes)
        logger.info(f"Pre-write checksum: {pre_checksum}")
        
        logger.info("Calculating post-write checksum...")
        post_checksum = sha256_file(checksum_target, read_bytes=checksum_bytes)
        logger.info(f"Post-write checksum: {post_checksum}")
        
        # Compare checksums
        checksum_match = pre_checksum == post_checksum
        if checksum_match:
            logger.info("SUCCESS: Checksums match!")
            status = "success"
        else:
            logger.warning("WARNING: Checksums do not match!")
            status = "checksum_mismatch"
        
        # Collect metadata
        end_time = time.time()
        metadata = {
            "target": target,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "unix_timestamp": int(time.time()),
            "duration_seconds": round(end_time - start_time, 2),
            "bytes_written": bytes_written,
            "pattern": pattern,
            "method": method,
            "pre_checksum": pre_checksum,
            "post_checksum": post_checksum,
            "checksum_match": checksum_match,
            "status": status,
            "test_file_path": file_path
        }
        
        # Save metadata
        try:
            with open(metadata_file, "w") as mf:
                json.dump(metadata, mf, indent=2)
            logger.info(f"Metadata saved to {metadata_file}")
            metadata["metadata_file"] = metadata_file
        except (OSError, IOError) as e:
            logger.error(f"Failed to save metadata to {metadata_file}: {e}")
            metadata["metadata_save_error"] = str(e)
        
        return metadata
        
    except TestWriteError as e:
        logger.error(f"Test write failed: {e}")
        error_metadata = {
            "target": target,
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "status": "error",
            "error": str(e),
            "pattern": pattern,
            "method": "unknown"
        }
        return error_metadata

def main():
    """Main function for standalone execution"""
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Write test data to device or mount point for wiping verification'
    )
    parser.add_argument('target', help='Device path (/dev/sdX) or mount point (/mnt)')
    parser.add_argument('--pattern', default="SUPER_SECRET_DATA_123456789",
                       help='Data pattern to write (default: SUPER_SECRET_DATA_123456789)')
    parser.add_argument('--output-dir', help='Directory for temporary and metadata files')
    parser.add_argument('--metadata-file', help='Path to save metadata JSON')
    parser.add_argument('--json', action='store_true', help='Output result as JSON')
    parser.add_argument('--quiet', '-q', action='store_true', help='Suppress info messages')
    
    args = parser.parse_args()
    
    # Adjust logging level
    if args.quiet:
        logging.getLogger().setLevel(logging.WARNING)
    
    # Check if we need root permissions
    if not is_mount_point(args.target) and os.geteuid() != 0:
        logger.error("Root permissions required for writing to raw devices. Use sudo.")
        return 1
    
    try:
        result = test_write_data(
            target=args.target,
            pattern=args.pattern,
            output_dir=args.output_dir,
            metadata_file=args.metadata_file
        )
        
        if args.json:
            print(json.dumps(result, indent=2))
        else:
            status = result.get('status', 'unknown')
            if status == 'success':
                print(f"✓ Test write completed successfully")
                print(f"  Target: {result['target']}")
                print(f"  Method: {result['method']}")
                print(f"  Bytes written: {result['bytes_written']}")
                print(f"  Duration: {result['duration_seconds']}s")
                print(f"  Checksums match: {result['checksum_match']}")
            elif status == 'error':
                print(f"✗ Test write failed: {result['error']}")
                return 1
            else:
                print(f"⚠ Test write completed with issues: {status}")
                return 2
        
        return 0
        
    except KeyboardInterrupt:
        logger.info("Test write interrupted by user")
        return 1
    except Exception as e:
        logger.error(f"Unexpected error: {e}")
        return 1

if __name__ == "__main__":
    sys.exit(main())