scripts/wg_api.py

108 lines
3.2 KiB
Python

import subprocess
import base64
from fastapi import FastAPI, Form, HTTPException, Query, Header, Depends
from fastapi.responses import JSONResponse
from pathlib import Path
import uvicorn
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
API_KEY_ENV = os.getenv("API_KEY")
WG_DIR = Path("/etc/wireguard")
SCRIPT_PATH = Path("/etc/wireguard/wg_config.sh")
app = FastAPI()
# Dependency to enforce API key
def verify_api_key(x_api_key: str = Header(...)):
if x_api_key != API_KEY_ENV:
raise HTTPException(status_code=403, detail="Invalid API key")
def is_valid_name(name):
return name.isalnum()
@app.post("/vpn")
async def create_vpn_client(
new: str = Form(...),
auth: None = Depends(verify_api_key)
):
client_name = new.strip()
if not is_valid_name(client_name):
raise HTTPException(status_code=400, detail="Invalid client name")
client_dir = WG_DIR / client_name
conf_file = client_dir / f"{client_name}.conf"
png_file = client_dir / f"{client_name}.png"
if client_dir.exists():
return JSONResponse(status_code=409, content={"error": "Client already exists"})
try:
subprocess.run(
["bash", str(SCRIPT_PATH), "add", client_name],
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
if not conf_file.exists() or not png_file.exists():
raise FileNotFoundError("Missing config or QR code after creation")
config_data = conf_file.read_text()
with open(png_file, "rb") as f:
qr_base64 = base64.b64encode(f.read()).decode("utf-8")
return {
"client": client_name,
"config": config_data,
"qr_base64": qr_base64
}
except subprocess.CalledProcessError as e:
raise HTTPException(status_code=500, detail=f"Script error: {e.stderr.decode()}")
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/vpn")
async def remove_vpn_client(
remove: str = Query(...),
auth: None = Depends(verify_api_key)
):
client_name = remove.strip()
if not is_valid_name(client_name):
raise HTTPException(status_code=400, detail="Invalid client name")
client_dir = WG_DIR / client_name
if not client_dir.exists():
raise HTTPException(status_code=404, detail="Client not found")
try:
subprocess.run(
["bash", str(SCRIPT_PATH), "remove", client_name],
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
return {"message": f"Client '{client_name}' removed."}
except subprocess.CalledProcessError as e:
raise HTTPException(status_code=500, detail=f"Script error: {e.stderr.decode()}")
@app.get("/vpn/list")
async def list_vpn_clients(auth: None = Depends(verify_api_key)):
clients = []
for dir in WG_DIR.iterdir():
if dir.is_dir():
conf = dir / f"{dir.name}.conf"
if conf.exists():
clients.append(dir.name)
return {"clients": clients}
if __name__ == "__main__":
uvicorn.run("wg_api:app", host="0.0.0.0", port=3000)