#!/usr/bin/env python3
"""
Generate WireGuard client configs from pfSense server config.
Generates new keypairs for each peer and outputs the public keys for pfSense update.
"""
import subprocess
import sys

def generate_keypair():
    """Generate a new WireGuard private/public keypair."""
    try:
        # Generate private key
        privkey = subprocess.check_output(['wg', 'genkey'], text=True).strip()
        # Derive public key
        pubkey = subprocess.check_output(
            ['wg', 'pubkey'],
            input=privkey,
            text=True
        ).strip()
        return privkey, pubkey
    except FileNotFoundError:
        print("Error: 'wg' command not found. Install wireguard-tools.", file=sys.stderr)
        sys.exit(1)


def parse_server_config(config_path):
    """Parse pfSense WireGuard server config file."""
    config = {}
    peers = []
    current_peer = None

    with open(config_path, 'r') as f:
        for line in f:
            line = line.strip()

            if line.startswith('[Interface]'):
                current_peer = None
            elif line.startswith('[Peer]'):
                current_peer = {}
                peers.append(current_peer)
            elif '=' in line:
                key, value = [x.strip() for x in line.split('=', 1)]

                if current_peer is None:
                    # Server interface section
                    config[key] = value
                else:
                    # Peer section
                    current_peer[key] = value

    return config, peers


def generate_client_config(server_config, peer, server_endpoint, new_privkey):
    """Generate client config for a peer."""
    # Extract server's public key
    server_pubkey = server_config.get('PrivateKey', '')
    if server_pubkey:
        # Derive public key from server's private key
        try:
            server_pubkey = subprocess.check_output(
                ['wg', 'pubkey'],
                input=server_pubkey,
                text=True
            ).strip()
        except:
            server_pubkey = ''

    # Get server's listen port
    listen_port = server_config.get('ListenPort', '51820')

    # Get peer's allowed IPs (this becomes the client's address)
    allowed_ips = peer.get('AllowedIPs', '')

    # Parse client IP from AllowedIPs (e.g., "10.0.0.2/32" -> "10.0.0.2/32")
    client_address = allowed_ips if allowed_ips else '10.0.0.2/24'

    # Build client config
    config = f"""[Interface]
PrivateKey = {new_privkey}
Address = {client_address}
DNS = 1.1.1.1

[Peer]
PublicKey = {server_pubkey}
Endpoint = {server_endpoint}:{listen_port}
AllowedIPs = 0.0.0.0/0
PersistentKeepalive = 25"""

    return config


def main():
    if len(sys.argv) < 3:
        print("Usage: python script.py <server_config_file> <server_endpoint_ip>")
        print("Example: python script.py wg0.conf 49.253.81.77")
        sys.exit(1)
    conf_file = sys.argv[1]
    server_endpoint = sys.argv[2]

    # for when I'm extra lazy
    # server_endpoint = "ChangeMeToYourWANIP"
    # for root, dirs, files in os.walk(os.getcwd()):
    #     for f in files:
    #         if f.endswith('.conf'):
    #             conf_files.append(os.path.join(root, f))

    try:
        with open(conf_file, "r"):
            pass
    except FileNotFoundError:
        print(f"Error: Config file not found.", file=sys.stderr)
        sys.exit(1)

    # Parse server config
    server_config, peers = parse_server_config(conf_file)

    if not peers:
        print("No peers found in config file.", file=sys.stderr)
        sys.exit(1)

    print(f"Found {len(peers)} peer(s) in config.\n")

    # Generate configs for each peer
    pfsense_updates = []

    for i, peer in enumerate(peers, 1):
        # Generate new keypair
        new_privkey, new_pubkey = generate_keypair()

        # Generate client config
        client_config = generate_client_config(
            server_config,
            peer,
            server_endpoint,
            new_privkey
        )

        # Save client config
        peer_id = peer.get('AllowedIPs', f'peer{i}').replace('/', '_')
        output_file = f"client_{peer_id}.conf"

        with open(output_file, 'w') as f:
            f.write(client_config)

        print(f"Generated: {output_file}")
        print("=" * 70)
        print(client_config)
        print("=" * 70)
        print()

        # Store for pfSense update instructions
        pfsense_updates.append({
            'old_pubkey': peer.get('PublicKey', 'N/A'),
            'new_pubkey': new_pubkey,
            'allowed_ips': peer.get('AllowedIPs', 'N/A')
        })

    # Print pfSense update instructions
    print("\n" + "=" * 70)
    print("PFSENSE UPDATE INSTRUCTIONS")
    print("=" * 70)
    print("\nIn pfSense, go to VPN > WireGuard > Peers and update each peer:\n")

    for i, update in enumerate(pfsense_updates, 1):
        print(f"Peer {i} (Allowed IPs: {update['allowed_ips']}):")
        print(f"  Old Public Key: {update['old_pubkey']}")
        print(f"  New Public Key: {update['new_pubkey']}")
        print()


if __name__ == '__main__':
    main()