Update generate_wireguard_config.py

generate_wireguard_config.py

@@ -1,81 +1,194 @@
-import os
-import subprocess
-
-WG_DIR = "/etc/wireguard"
-SERVER_PRIVATE = os.path.join(WG_DIR, "privatekey")
-SERVER_PUBLIC = os.path.join(WG_DIR, "publickey")
-PEER_PRIVATE = os.path.join(WG_DIR, "peer1_privatekey")
-PEER_PUBLIC = os.path.join(WG_DIR, "peer1_publickey")
-
-
-def generate_keypair(private_path, public_path):
-    priv = subprocess.check_output(["wg", "genkey"]).decode().strip()
-    with open(private_path, "w") as f:
-        f.write(priv)
-    pub = subprocess.check_output(["wg", "pubkey"], input=priv.encode()).decode().strip()
-    with open(public_path, "w") as f:
-        f.write(pub)
-    return priv, pub
-
-def main():
-    os.makedirs(WG_DIR, exist_ok=True)
-    # Generate keys
-    server_priv, server_pub = generate_keypair(SERVER_PRIVATE, SERVER_PUBLIC)
-    peer_priv, peer_pub = generate_keypair(PEER_PRIVATE, PEER_PUBLIC)
-
-    # Server config
-    wg0_conf = f"""
-[Interface]
-Address = 10.10.0.1/24
-ListenPort = 51820
-PrivateKey = {server_priv}
-PostUp = iptables -t nat -A POSTROUTING -o eth0 -j MASQUERADE
-PostDown = iptables -t nat -D POSTROUTING -o eth0 -j MASQUERADE
-
-[Peer]
-PublicKey = {peer_pub}
-AllowedIPs = 10.10.0.2/32
-"""
-    with open(os.path.join(WG_DIR, "wg0.conf"), "w") as f:
-        f.write(wg0_conf.strip())
-
-    # Try to auto-detect public IP
-    import socket
-    import urllib.request
-    public_ip = None
-    try:
-        # Try to get public IP from external service
-        with urllib.request.urlopen('https://api.ipify.org') as response:
-            public_ip = response.read().decode().strip()
-    except Exception:
-        # Fallback: get IP of eth0
-        try:
-            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-            s.connect(("8.8.8.8", 80))
-            public_ip = s.getsockname()[0]
-            s.close()
-        except Exception:
-            public_ip = "<UNKNOWN>"
-
-    print(f"[INFO] Detected public IP: {public_ip}")
-
-    # Peer config
-    peer_conf = f"""
-[Interface]
-PrivateKey = {peer_priv}
-Address = 10.10.0.2/24
-DNS = 1.1.1.1
-
-[Peer]
-PublicKey = {server_pub}
-Endpoint = {public_ip}:51820
-AllowedIPs = 0.0.0.0/0
-PersistentKeepalive = 25
-"""
-    with open(os.path.join(WG_DIR, "peer1.conf"), "w") as f:
-        f.write(peer_conf.strip())
-
-    print("✔ WireGuard config generated.")
-
-if __name__ == "__main__":
-    main()
+import os
+import subprocess
+import sys
+
+WG_DIR = "/etc/wireguard"
+SERVER_PRIVATE = os.path.join(WG_DIR, "privatekey")
+SERVER_PUBLIC = os.path.join(WG_DIR, "publickey")
+PEER_PRIVATE = os.path.join(WG_DIR, "peer1_privatekey")
+PEER_PUBLIC = os.path.join(WG_DIR, "peer1_publickey")
+
+
+def check_root():
+    """Check if running as root"""
+    if os.geteuid() != 0:
+        print("❌ This script must be run as root (use sudo)")
+        sys.exit(1)
+
+
+def check_wireguard():
+    """Check if WireGuard is available"""
+    try:
+        subprocess.run(["wg", "--version"], check=True, capture_output=True)
+        print("✔ WireGuard tools found")
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        print("❌ WireGuard tools not found. Install with:")
+        print("   Amazon Linux: sudo yum install wireguard-tools")
+        print("   Ubuntu/Debian: sudo apt install wireguard")
+        sys.exit(1)
+
+
+def load_wireguard_module():
+    """Try to load WireGuard kernel module"""
+    try:
+        subprocess.run(["modprobe", "wireguard"], check=True, capture_output=True)
+        print("✔ WireGuard kernel module loaded")
+    except subprocess.CalledProcessError:
+        print("⚠ Warning: Could not load WireGuard kernel module")
+        print("  This might be a container/virtualized environment")
+        print("  WireGuard may still work with userspace implementation")
+
+
+def enable_ip_forwarding():
+    """Enable IP forwarding if possible"""
+    try:
+        # Try to enable IP forwarding
+        subprocess.run(["sysctl", "-w", "net.ipv4.ip_forward=1"], 
+                      check=True, capture_output=True)
+        subprocess.run(["sysctl", "-w", "net.ipv6.conf.all.forwarding=1"], 
+                      check=True, capture_output=True)
+        print("✔ IP forwarding enabled")
+    except subprocess.CalledProcessError:
+        print("⚠ Warning: Could not enable IP forwarding")
+        print("  This might be a read-only filesystem or container")
+        print("  You may need to enable it manually or in the host system")
+
+
+def generate_keypair(private_path, public_path):
+    """Generate WireGuard key pair"""
+    try:
+        priv = subprocess.check_output(["wg", "genkey"]).decode().strip()
+        with open(private_path, "w") as f:
+            f.write(priv)
+        os.chmod(private_path, 0o600)  # Secure permissions
+        
+        pub = subprocess.check_output(["wg", "pubkey"], input=priv.encode()).decode().strip()
+        with open(public_path, "w") as f:
+            f.write(pub)
+        os.chmod(public_path, 0o644)
+        
+        return priv, pub
+    except Exception as e:
+        print(f"❌ Error generating keypair: {e}")
+        sys.exit(1)
+
+
+def get_public_ip():
+    """Auto-detect public IP address"""
+    import socket
+    import urllib.request
+    
+    # Try external service first
+    try:
+        with urllib.request.urlopen('https://api.ipify.org', timeout=5) as response:
+            return response.read().decode().strip()
+    except Exception:
+        pass
+    
+    # Fallback to local interface IP
+    try:
+        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+        s.connect(("8.8.8.8", 80))
+        ip = s.getsockname()[0]
+        s.close()
+        return ip
+    except Exception:
+        return "<REPLACE_WITH_YOUR_PUBLIC_IP>"
+
+
+def create_server_config(server_priv, peer_pub, interface="eth0"):
+    """Create server configuration"""
+    return f"""[Interface]
+Address = 10.10.0.1/24
+ListenPort = 51820
+PrivateKey = {server_priv}
+PostUp = iptables -t nat -A POSTROUTING -o {interface} -j MASQUERADE
+PostUp = iptables -A FORWARD -i wg0 -j ACCEPT
+PostUp = iptables -A FORWARD -o wg0 -j ACCEPT
+PostDown = iptables -t nat -D POSTROUTING -o {interface} -j MASQUERADE
+PostDown = iptables -D FORWARD -i wg0 -j ACCEPT
+PostDown = iptables -D FORWARD -o wg0 -j ACCEPT
+
+[Peer]
+PublicKey = {peer_pub}
+AllowedIPs = 10.10.0.2/32"""
+
+
+def create_peer_config(peer_priv, server_pub, public_ip):
+    """Create peer configuration"""
+    return f"""[Interface]
+PrivateKey = {peer_priv}
+Address = 10.10.0.2/24
+DNS = 1.1.1.1, 8.8.8.8
+
+[Peer]
+PublicKey = {server_pub}
+Endpoint = {public_ip}:51820
+AllowedIPs = 0.0.0.0/0
+PersistentKeepalive = 25"""
+
+
+def main():
+    print("🔧 WireGuard Configuration Generator")
+    print("=" * 40)
+    
+    # Pre-flight checks
+    check_root()
+    check_wireguard()
+    load_wireguard_module()
+    enable_ip_forwarding()
+    
+    # Create directory
+    try:
+        os.makedirs(WG_DIR, exist_ok=True)
+        print(f"✔ Created directory: {WG_DIR}")
+    except Exception as e:
+        print(f"❌ Could not create directory {WG_DIR}: {e}")
+        sys.exit(1)
+    
+    # Generate keys
+    print("\n🔑 Generating keypairs...")
+    server_priv, server_pub = generate_keypair(SERVER_PRIVATE, SERVER_PUBLIC)
+    peer_priv, peer_pub = generate_keypair(PEER_PRIVATE, PEER_PUBLIC)
+    print("✔ Server and peer keypairs generated")
+    
+    # Get public IP
+    public_ip = get_public_ip()
+    print(f"🌐 Detected public IP: {public_ip}")
+    
+    # Create server config
+    server_config = create_server_config(server_priv, peer_pub)
+    try:
+        with open(os.path.join(WG_DIR, "wg0.conf"), "w") as f:
+            f.write(server_config)
+        os.chmod(os.path.join(WG_DIR, "wg0.conf"), 0o600)
+        print("✔ Server config written to wg0.conf")
+    except Exception as e:
+        print(f"❌ Error writing server config: {e}")
+        sys.exit(1)
+    
+    # Create peer config
+    peer_config = create_peer_config(peer_priv, server_pub, public_ip)
+    try:
+        with open(os.path.join(WG_DIR, "peer1.conf"), "w") as f:
+            f.write(peer_config)
+        os.chmod(os.path.join(WG_DIR, "peer1.conf"), 0o600)
+        print("✔ Peer config written to peer1.conf")
+    except Exception as e:
+        print(f"❌ Error writing peer config: {e}")
+        sys.exit(1)
+    
+    print("\n🎉 WireGuard configuration completed!")
+    print("\nNext steps:")
+    print("1. Start WireGuard: sudo wg-quick up wg0")
+    print("2. Enable at boot: sudo systemctl enable wg-quick@wg0")
+    print("3. Copy peer1.conf to your client device")
+    print("4. Check status: sudo wg show")
+    
+    if public_ip == "<REPLACE_WITH_YOUR_PUBLIC_IP>":
+        print("\n⚠ Warning: Could not detect public IP!")
+        print("  Edit peer1.conf and replace the placeholder with your server's public IP")
+
+
+if __name__ == "__main__":
+    main()