Files
trade-iac/scripts/wireguard/wg_serverctl.py

420 lines
12 KiB
Python
Executable File

#!/usr/bin/env python3
import argparse
import os
import re
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Optional, List
WG_KEY_RE = re.compile(r"^[A-Za-z0-9+/]{43}=$")
@dataclass
class Cfg:
wg_dir: Path = Path("/etc/wireguard")
clients_dir: Path = Path("/etc/wireguard/clients")
server_key: Path = Path("/etc/wireguard/server.key")
server_pub: Path = Path("/etc/wireguard/server.pub")
wg0_conf: Path = Path("/etc/wireguard/wg0.conf")
wg_addr: str = "10.66.66.1/24"
wg_port: int = 51820
keepalive: int = 25
client_net_prefix: str = "10.66.66."
client_ip_start: int = 2
client_ip_end: int = 254
def require_root():
if os.geteuid() != 0:
print("ERROR: run as root (sudo).", file=sys.stderr)
sys.exit(1)
def run(
cmd: List[str],
capture: bool = False,
check: bool = True,
input_text: Optional[str] = None,
) -> subprocess.CompletedProcess:
return subprocess.run(
cmd,
text=True,
input=input_text,
stdout=subprocess.PIPE if capture else None,
stderr=subprocess.PIPE if capture else None,
check=check,
)
def ensure_dirs(cfg: Cfg):
cfg.wg_dir.mkdir(parents=True, exist_ok=True)
cfg.clients_dir.mkdir(parents=True, exist_ok=True)
os.chmod(cfg.wg_dir, 0o700)
os.chmod(cfg.clients_dir, 0o700)
def read_first_nonempty_line(p: Path) -> str:
if not p.exists():
return ""
for line in p.read_text(encoding="utf-8", errors="ignore").splitlines():
line = line.strip().replace("\r", "")
if line:
return line
return ""
def validate_key(k: str, what: str):
if not WG_KEY_RE.match(k):
raise ValueError(f"{what} has invalid WireGuard key format: {k!r}")
def chmod600_root(p: Path):
os.chmod(p, 0o600)
os.chown(p, 0, 0)
def backup_if_exists(p: Path):
if p.exists():
ts = datetime.now().strftime("%Y-%m-%d_%H%M%S")
b = p.with_name(p.name + f".bak.{ts}")
shutil.copy2(p, b)
def gen_priv_key_to(path: Path):
priv = run(["wg", "genkey"], capture=True).stdout.strip()
validate_key(priv, "generated private key")
path.write_text(priv + "\n", encoding="utf-8")
chmod600_root(path)
def pubkey_from_priv(priv: str) -> str:
validate_key(priv, "private key")
pub = run(["wg", "pubkey"], capture=True, input_text=priv + "\n").stdout.strip()
validate_key(pub, "generated public key")
return pub
def detect_public_ip() -> str:
# Works on Ubuntu/Debian in typical setups
out = run(
[
"bash",
"-lc",
"ip route get 1.1.1.1 | awk '{for (i=1;i<=NF;i++) if ($i==\"src\") {print $(i+1); exit}}'",
],
capture=True,
).stdout.strip()
if not out:
raise RuntimeError("Could not detect server public IP.")
return out
def restart_wg(unit: str = "wg-quick@wg0"):
run(["systemctl", "restart", unit], capture=False, check=True)
# optional status (non-fatal)
try:
run(["systemctl", "--no-pager", "-l", "status", unit], capture=False, check=False)
except Exception:
pass
def list_clients(cfg: Cfg) -> List[str]:
pubs = sorted(cfg.clients_dir.glob("*.pub"))
return [p.stem for p in pubs]
def used_octets(cfg: Cfg) -> set:
s = set()
for f in cfg.clients_dir.glob("*.ip"):
v = read_first_nonempty_line(f)
if v.isdigit():
s.add(int(v))
return s
def alloc_octet(cfg: Cfg) -> int:
used = used_octets(cfg)
for i in range(cfg.client_ip_start, cfg.client_ip_end + 1):
if i not in used:
return i
raise RuntimeError("No free IPs left in pool.")
def gen_server(cfg: Cfg, force: bool):
ensure_dirs(cfg)
if not force and (cfg.server_key.exists() or cfg.server_pub.exists()):
raise RuntimeError("Server key already exists. Use --force to overwrite.")
gen_priv_key_to(cfg.server_key)
priv = read_first_nonempty_line(cfg.server_key)
pub = pubkey_from_priv(priv)
cfg.server_pub.write_text(pub + "\n", encoding="utf-8")
chmod600_root(cfg.server_pub)
print("OK: generated server keys:")
print(f" - {cfg.server_key} (PRIVATE)")
print(f" - {cfg.server_pub} (PUBLIC)")
print("Server public key:")
print(pub)
def gen_client(cfg: Cfg, name: str, force: bool):
ensure_dirs(cfg)
keyf = cfg.clients_dir / f"{name}.key"
pubf = cfg.clients_dir / f"{name}.pub"
if not force and (keyf.exists() or pubf.exists()):
raise RuntimeError(f"Client {name} exists. Use --force to overwrite.")
gen_priv_key_to(keyf)
priv = read_first_nonempty_line(keyf)
pub = pubkey_from_priv(priv)
pubf.write_text(pub + "\n", encoding="utf-8")
chmod600_root(pubf)
print("OK: generated client keys:")
print(f" - {keyf} (PRIVATE)")
print(f" - {pubf} (PUBLIC)")
print("Client public key:")
print(pub)
def generate_wg0_conf(cfg: Cfg):
ensure_dirs(cfg)
if not cfg.server_key.exists():
raise RuntimeError("Missing server.key. Run gen-server first.")
server_priv = read_first_nonempty_line(cfg.server_key)
validate_key(server_priv, "server private key")
backup_if_exists(cfg.wg0_conf)
lines: List[str] = []
lines += [
"[Interface]",
f"Address = {cfg.wg_addr}",
f"ListenPort = {cfg.wg_port}",
f"PrivateKey = {server_priv}",
]
for name in sorted(list_clients(cfg)):
pubf = cfg.clients_dir / f"{name}.pub"
ipf = cfg.clients_dir / f"{name}.ip"
if not ipf.exists():
raise RuntimeError(f"Client {name} has no .ip file (run add {name}).")
pub = read_first_nonempty_line(pubf)
validate_key(pub, f"client {name} pubkey")
octet = read_first_nonempty_line(ipf)
if not octet.isdigit():
raise RuntimeError(f"Invalid ip octet in {ipf}: {octet!r}")
o = int(octet)
if o < 2 or o > 254:
raise RuntimeError(f"Invalid ip octet in {ipf}: {o}")
lines += [
"",
"[Peer]",
f"# {name}",
f"PublicKey = {pub}",
f"AllowedIPs = {cfg.client_net_prefix}{o}/32",
f"PersistentKeepalive = {cfg.keepalive}",
]
cfg.wg0_conf.write_text("\n".join(lines) + "\n", encoding="utf-8")
chmod600_root(cfg.wg0_conf)
print(f"OK: generated {cfg.wg0_conf}")
def add_client(cfg: Cfg, name: str, pubkey: Optional[str], ip: Optional[str], no_restart: bool):
ensure_dirs(cfg)
pubf = cfg.clients_dir / f"{name}.pub"
ipf = cfg.clients_dir / f"{name}.ip"
# resolve pubkey
if pubkey:
pubkey = pubkey.strip()
validate_key(pubkey, "provided client public key")
pubf.write_text(pubkey + "\n", encoding="utf-8")
chmod600_root(pubf)
else:
if not pubf.exists():
raise RuntimeError(f"Missing {pubf}. Run gen-client {name} or provide --pubkey.")
pubkey = read_first_nonempty_line(pubf)
validate_key(pubkey, f"client {name} pubkey")
# resolve ip octet
if ip is None:
octet = alloc_octet(cfg)
else:
ip = ip.strip()
if ip.startswith(cfg.client_net_prefix):
ip = ip[len(cfg.client_net_prefix) :]
if not ip.isdigit():
raise RuntimeError("--ip must be an octet (2..254) or full IP like 10.66.66.X")
octet = int(ip)
if octet < 2 or octet > 254:
raise RuntimeError("--ip must be 2..254")
if octet in used_octets(cfg):
raise RuntimeError(f"IP octet {octet} already used")
ipf.write_text(f"{octet}\n", encoding="utf-8")
chmod600_root(ipf)
print(f"OK: added peer {name} -> {cfg.client_net_prefix}{octet}/32")
generate_wg0_conf(cfg)
if not no_restart:
restart_wg()
def del_client(cfg: Cfg, name: str, no_restart: bool):
ensure_dirs(cfg)
pubf = cfg.clients_dir / f"{name}.pub"
ipf = cfg.clients_dir / f"{name}.ip"
# keep .key on purpose
if pubf.exists():
pubf.unlink()
if ipf.exists():
ipf.unlink()
print(f"OK: removed peer {name} (pub/ip). Private key kept if existed.")
generate_wg0_conf(cfg)
if not no_restart:
restart_wg()
def export_client_conf(cfg: Cfg, name: str, endpoint: Optional[str]):
ensure_dirs(cfg)
keyf = cfg.clients_dir / f"{name}.key"
ipf = cfg.clients_dir / f"{name}.ip"
if not keyf.exists():
raise RuntimeError(f"Missing {keyf}. Run gen-client {name} on server or create key on client.")
if not ipf.exists():
raise RuntimeError(f"Missing {ipf}. Run add {name} first.")
if not cfg.server_pub.exists():
raise RuntimeError("Missing server.pub. Run gen-server first.")
client_priv = read_first_nonempty_line(keyf)
validate_key(client_priv, f"client {name} private key")
server_pub = read_first_nonempty_line(cfg.server_pub)
validate_key(server_pub, "server public key")
octet = int(read_first_nonempty_line(ipf))
if endpoint is None:
endpoint = f"{detect_public_ip()}:{cfg.wg_port}"
out = cfg.clients_dir / f"{name}.conf"
content = "\n".join(
[
"[Interface]",
f"Address = {cfg.client_net_prefix}{octet}/32",
f"PrivateKey = {client_priv}",
"",
"[Peer]",
f"PublicKey = {server_pub}",
f"Endpoint = {endpoint}",
# Default: only reach the server's WG IP via tunnel.
"AllowedIPs = 10.66.66.1/32",
f"PersistentKeepalive = {cfg.keepalive}",
"",
]
)
out.write_text(content, encoding="utf-8")
chmod600_root(out)
print(f"OK: wrote {out}")
def cmd_list(cfg: Cfg):
ensure_dirs(cfg)
names = list_clients(cfg)
if not names:
print("(no clients)")
return
for n in names:
ipf = cfg.clients_dir / f"{n}.ip"
octet = read_first_nonempty_line(ipf) if ipf.exists() else "?"
print(f"{n}\t{cfg.client_net_prefix}{octet}/32")
def cmd_show(cfg: Cfg, name: str):
ensure_dirs(cfg)
pubf = cfg.clients_dir / f"{name}.pub"
ipf = cfg.clients_dir / f"{name}.ip"
keyf = cfg.clients_dir / f"{name}.key"
print(f"Name: {name}")
if ipf.exists():
print(f"IP: {cfg.client_net_prefix}{read_first_nonempty_line(ipf)}/32")
else:
print("IP: (not assigned)")
if pubf.exists():
print(f"PUB: {read_first_nonempty_line(pubf)}")
else:
print("PUB: (missing)")
print(f"KEY: {'exists' if keyf.exists() else 'missing'} ({keyf})")
def main():
require_root()
cfg = Cfg()
p = argparse.ArgumentParser(prog="wg_serverctl.py")
sub = p.add_subparsers(dest="cmd", required=True)
s = sub.add_parser("gen-server")
s.add_argument("--force", action="store_true")
g = sub.add_parser("gen-client")
g.add_argument("name")
g.add_argument("--force", action="store_true")
a = sub.add_parser("add")
a.add_argument("name")
a.add_argument("--pubkey", help="Client public key (base64, ends with =)")
a.add_argument("--ip", help="IP octet (2..254) or full 10.66.66.X")
a.add_argument("--no-restart", action="store_true")
d = sub.add_parser("del")
d.add_argument("name")
d.add_argument("--no-restart", action="store_true")
ap = sub.add_parser("apply")
ap.add_argument("--no-restart", action="store_true")
sub.add_parser("list")
sh = sub.add_parser("show")
sh.add_argument("name")
ex = sub.add_parser("export")
ex.add_argument("name")
ex.add_argument("--endpoint", help="Override endpoint ip:port (default: auto-detect public ip)")
args = p.parse_args()
try:
if args.cmd == "gen-server":
gen_server(cfg, args.force)
elif args.cmd == "gen-client":
gen_client(cfg, args.name, args.force)
elif args.cmd == "add":
add_client(cfg, args.name, args.pubkey, args.ip, args.no_restart)
elif args.cmd == "del":
del_client(cfg, args.name, args.no_restart)
elif args.cmd == "apply":
generate_wg0_conf(cfg)
if not args.no_restart:
restart_wg()
elif args.cmd == "list":
cmd_list(cfg)
elif args.cmd == "show":
cmd_show(cfg, args.name)
elif args.cmd == "export":
export_client_conf(cfg, args.name, args.endpoint)
else:
p.print_help()
except Exception as e:
print(f"ERROR: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()