Peer enhancements: notes, traffic quotas with auto-disable, batch create, usage tracking; relay subnets for bridging into other tunnels

Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>

app/db.py

@@ -1,4 +1,4 @@
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, text
 from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
 
 from .config import settings
@@ -21,3 +21,24 @@ def get_db():
         yield db
     finally:
         db.close()
+
+
+PEER_MIGRATIONS = {
+    "note": "ALTER TABLE peers ADD COLUMN note VARCHAR(256) NOT NULL DEFAULT ''",
+    "quota_bytes": "ALTER TABLE peers ADD COLUMN quota_bytes INTEGER NOT NULL DEFAULT 0",
+    "cum_rx": "ALTER TABLE peers ADD COLUMN cum_rx INTEGER NOT NULL DEFAULT 0",
+    "cum_tx": "ALTER TABLE peers ADD COLUMN cum_tx INTEGER NOT NULL DEFAULT 0",
+    "last_rx": "ALTER TABLE peers ADD COLUMN last_rx INTEGER NOT NULL DEFAULT 0",
+    "last_tx": "ALTER TABLE peers ADD COLUMN last_tx INTEGER NOT NULL DEFAULT 0",
+}
+
+
+def run_migrations() -> None:
+    with engine.connect() as conn:
+        existing = {
+            row[1] for row in conn.execute(text("PRAGMA table_info(peers)"))
+        }
+        for column, ddl in PEER_MIGRATIONS.items():
+            if existing and column not in existing:
+                conn.execute(text(ddl))
+        conn.commit()

app/main.py

@@ -13,9 +13,9 @@ from sqlalchemy.orm import Session
 from . import service, wireguard
 from .auth import SESSION_COOKIE, create_session_token, logged_in, verify_credentials
 from .config import settings
-from .db import Base, SessionLocal, engine, get_db
+from .db import Base, SessionLocal, engine, get_db, run_migrations
 
-TRAFFIC_SAMPLE_INTERVAL = 300
+TRAFFIC_SAMPLE_INTERVAL = 60
 
 templates = Jinja2Templates(directory="app/templates")
 
@@ -25,6 +25,7 @@ async def _background_sampler() -> None:
         await asyncio.sleep(TRAFFIC_SAMPLE_INTERVAL)
         db = SessionLocal()
         try:
+            service.accumulate_usage(db)
             service.disable_expired_peers(db)
             service.sample_traffic(db)
         except Exception:
@@ -36,6 +37,7 @@ async def _background_sampler() -> None:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     Base.metadata.create_all(engine)
+    run_migrations()
     db = SessionLocal()
     try:
         service.apply_config(db)
@@ -114,14 +116,26 @@ def peers_page(request: Request, db: Session = Depends(get_db)):
     )
 
 
+def _parse_quota_gib(value: str) -> int:
+    quota = float(value) if value.strip() else 0
+    return int(quota * 1024**3)
+
+
 @app.post("/peers", dependencies=[logged_in])
 def create_peer(
     db: Session = Depends(get_db),
     name: str = Form(...),
     expires_at: str = Form(""),
+    note: str = Form(""),
+    quota_gib: str = Form(""),
+    count: int = Form(1),
 ):
     expiry = datetime.fromisoformat(expires_at) if expires_at else None
-    peer = service.create_peer(db, name.strip(), expiry)
+    quota = _parse_quota_gib(quota_gib)
+    if count > 1:
+        service.create_peers_batch(db, name.strip(), min(count, 50), expiry, note.strip(), quota)
+        return RedirectResponse("/peers", status_code=303)
+    peer = service.create_peer(db, name.strip(), expiry, note.strip(), quota)
     return RedirectResponse(f"/peers/{peer.id}", status_code=303)
 
 
@@ -143,6 +157,24 @@ def peer_detail(request: Request, peer_id: int, db: Session = Depends(get_db)):
     )
 
 
+@app.post("/peers/{peer_id}/update", dependencies=[logged_in])
+def update_peer(
+    peer_id: int,
+    db: Session = Depends(get_db),
+    note: str = Form(""),
+    quota_gib: str = Form(""),
+):
+    peer = _get_peer_or_404(db, peer_id)
+    service.update_peer(db, peer, note.strip(), _parse_quota_gib(quota_gib))
+    return RedirectResponse(f"/peers/{peer_id}", status_code=303)
+
+
+@app.post("/peers/{peer_id}/reset-usage", dependencies=[logged_in])
+def reset_usage(peer_id: int, db: Session = Depends(get_db)):
+    service.reset_peer_usage(db, _get_peer_or_404(db, peer_id))
+    return RedirectResponse(f"/peers/{peer_id}", status_code=303)
+
+
 @app.post("/peers/{peer_id}/toggle", dependencies=[logged_in])
 def toggle_peer(peer_id: int, db: Session = Depends(get_db)):
     service.toggle_peer(db, _get_peer_or_404(db, peer_id))
@@ -211,6 +243,11 @@ def api_status(db: Session = Depends(get_db)):
                 "name": peer.name,
                 "address": peer.address,
                 "enabled": peer.enabled,
+                "note": peer.note,
+                "quota_bytes": peer.quota_bytes,
+                "cum_rx": peer.cum_rx,
+                "cum_tx": peer.cum_tx,
+                "over_quota": peer.over_quota,
                 "online": (ps := status.peers.get(peer.public_key)) is not None and ps.online,
                 "endpoint": ps.endpoint if ps else None,
                 "latest_handshake": ps.latest_handshake.isoformat()

app/models.py

@@ -22,6 +22,20 @@ class Peer(Base):
     enabled: Mapped[bool] = mapped_column(Boolean, default=True)
     expires_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
     created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    note: Mapped[str] = mapped_column(String(256), default="")
+    quota_bytes: Mapped[int] = mapped_column(Integer, default=0)
+    cum_rx: Mapped[int] = mapped_column(Integer, default=0)
+    cum_tx: Mapped[int] = mapped_column(Integer, default=0)
+    last_rx: Mapped[int] = mapped_column(Integer, default=0)
+    last_tx: Mapped[int] = mapped_column(Integer, default=0)
+
+    @property
+    def cum_total(self) -> int:
+        return self.cum_rx + self.cum_tx
+
+    @property
+    def over_quota(self) -> bool:
+        return self.quota_bytes > 0 and self.cum_total >= self.quota_bytes
 
 
 class TrafficSample(Base):

app/service.py

@@ -16,7 +16,13 @@ def get_peer(db: Session, peer_id: int) -> Peer | None:
     return db.get(Peer, peer_id)
 
 
-def create_peer(db: Session, name: str, expires_at: datetime | None = None) -> Peer:
+def create_peer(
+    db: Session,
+    name: str,
+    expires_at: datetime | None = None,
+    note: str = "",
+    quota_bytes: int = 0,
+) -> Peer:
     private, public = wireguard.generate_keypair()
     psk = wireguard.genpsk()
     address = wireguard.next_free_address([p.address for p in list_peers(db)])
@@ -27,6 +33,8 @@ def create_peer(db: Session, name: str, expires_at: datetime | None = None) -> P
         preshared_key_enc=encrypt(psk),
         address=address,
         expires_at=expires_at,
+        note=note,
+        quota_bytes=quota_bytes,
     )
     db.add(peer)
     db.commit()
@@ -34,8 +42,47 @@ def create_peer(db: Session, name: str, expires_at: datetime | None = None) -> P
     return peer
 
 
+def create_peers_batch(
+    db: Session,
+    base_name: str,
+    count: int,
+    expires_at: datetime | None = None,
+    note: str = "",
+    quota_bytes: int = 0,
+) -> list[Peer]:
+    existing = {p.name for p in list_peers(db)}
+    peers = []
+    index = 1
+    for _ in range(count):
+        while f"{base_name}-{index}" in existing:
+            index += 1
+        name = f"{base_name}-{index}"
+        existing.add(name)
+        peers.append(create_peer(db, name, expires_at, note, quota_bytes))
+    return peers
+
+
+def update_peer(db: Session, peer: Peer, note: str, quota_bytes: int) -> Peer:
+    peer.note = note
+    peer.quota_bytes = quota_bytes
+    db.commit()
+    if not peer.over_quota:
+        apply_config(db)
+    return peer
+
+
+def reset_peer_usage(db: Session, peer: Peer) -> Peer:
+    peer.cum_rx = 0
+    peer.cum_tx = 0
+    db.commit()
+    apply_config(db)
+    return peer
+
+
 def rotate_peer_keys(db: Session, peer: Peer) -> Peer:
     private, public = wireguard.generate_keypair()
+    peer.last_rx = 0
+    peer.last_tx = 0
     peer.public_key = public
     peer.private_key_enc = encrypt(private)
     peer.preshared_key_enc = encrypt(wireguard.genpsk())
@@ -61,10 +108,12 @@ def disable_expired_peers(db: Session) -> bool:
     now = datetime.now(timezone.utc)
     changed = False
     for peer in list_peers(db):
+        if not peer.enabled:
+            continue
         expires = peer.expires_at
         if expires is not None and expires.tzinfo is None:
             expires = expires.replace(tzinfo=timezone.utc)
-        if peer.enabled and expires is not None and expires <= now:
+        if (expires is not None and expires <= now) or peer.over_quota:
             peer.enabled = False
             changed = True
     if changed:
@@ -73,6 +122,27 @@ def disable_expired_peers(db: Session) -> bool:
     return changed
 
 
+def accumulate_usage(db: Session) -> None:
+    status = wireguard.get_status()
+    changed = False
+    for peer in list_peers(db):
+        peer_status = status.peers.get(peer.public_key)
+        if peer_status is None:
+            continue
+        rx, tx = peer_status.rx_bytes, peer_status.tx_bytes
+        # wg counters reset on interface restart or peer re-add
+        delta_rx = rx - peer.last_rx if rx >= peer.last_rx else rx
+        delta_tx = tx - peer.last_tx if tx >= peer.last_tx else tx
+        if delta_rx or delta_tx or rx != peer.last_rx or tx != peer.last_tx:
+            peer.cum_rx += delta_rx
+            peer.cum_tx += delta_tx
+            peer.last_rx = rx
+            peer.last_tx = tx
+            changed = True
+    if changed:
+        db.commit()
+
+
 def apply_config(db: Session) -> None:
     private, _ = wireguard.ensure_server_keys()
     wireguard.sync_peers(private, list_peers(db))

app/static/style.css

@@ -69,7 +69,10 @@ input {
   border-radius: 6px;
   font-size: 0.9rem;
 }
-.add-peer { display: flex; gap: 0.5rem; align-items: center; margin-bottom: 1rem; }
+.add-peer { display: flex; gap: 0.5rem; align-items: center; margin-bottom: 0.25rem; flex-wrap: wrap; }
+input.narrow { width: 80px; }
+form.stack label { display: block; margin-bottom: 0.5rem; }
+form.stack { margin-top: 0.75rem; }
 
 .login-page { display: flex; align-items: center; justify-content: center; min-height: 100vh; }
 .login-box {

app/templates/peer_detail.html

@@ -13,6 +13,19 @@
     </p>
     <p class="muted">Created {{ peer.created_at.strftime("%Y-%m-%d %H:%M") }}</p>
     {% if peer.expires_at %}<p class="muted">Expires {{ peer.expires_at.strftime("%Y-%m-%d %H:%M") }}</p>{% endif %}
+    <p>Usage: {{ peer.cum_total | fmt_bytes }}
+      {% if peer.quota_bytes %} / {{ peer.quota_bytes | fmt_bytes }}
+        {% if peer.over_quota %}<span class="badge err">over quota</span>{% endif %}
+      {% endif %}
+    </p>
+    <form method="post" action="/peers/{{ peer.id }}/update" class="stack">
+      <label class="muted">Note <input name="note" value="{{ peer.note }}"></label>
+      <label class="muted">Quota GiB
+        <input name="quota_gib" type="number" step="0.1" min="0" class="narrow"
+               value="{{ '%.1f' % (peer.quota_bytes / 1073741824) if peer.quota_bytes else '' }}" placeholder="0 = unlimited">
+      </label>
+      <button type="submit">Save</button>
+    </form>
   </div>
   <div class="card">
     <h3>Live status</h3>
@@ -40,6 +53,10 @@
         onsubmit="return confirm('Rotate keys? The old client config stops working immediately.')">
     <button type="submit">Rotate keys</button>
   </form>
+  <form method="post" action="/peers/{{ peer.id }}/reset-usage" class="inline"
+        onsubmit="return confirm('Reset usage counters to zero?')">
+    <button type="submit">Reset usage</button>
+  </form>
   <form method="post" action="/peers/{{ peer.id }}/delete" class="inline"
         onsubmit="return confirm('Delete peer {{ peer.name }}?')">
     <button type="submit" class="danger">Delete</button>

app/templates/peers.html

@@ -5,13 +5,17 @@
 
 <form method="post" action="/peers" class="add-peer">
   <input name="name" placeholder="Peer name" required>
+  <input name="note" placeholder="Note (optional)">
+  <label class="muted">Quota GiB <input name="quota_gib" type="number" step="0.1" min="0" class="narrow" placeholder="0"></label>
+  <label class="muted">Count <input name="count" type="number" min="1" max="50" value="1" class="narrow"></label>
   <label class="muted">Expires <input name="expires_at" type="datetime-local"></label>
-  <button type="submit">Add peer</button>
+  <button type="submit">Add peer(s)</button>
 </form>
+<p class="muted">Count &gt; 1 creates a batch named name-1, name-2, ... Quota 0 = unlimited.</p>
 
 <table>
   <thead>
-    <tr><th>Name</th><th>Address</th><th>Enabled</th><th>Status</th><th>Created</th><th></th></tr>
+    <tr><th>Name</th><th>Address</th><th>Enabled</th><th>Status</th><th>Usage</th><th>Note</th><th></th></tr>
   </thead>
   <tbody>
   {% for peer in peers %}
@@ -19,9 +23,10 @@
     <tr>
       <td><a href="/peers/{{ peer.id }}">{{ peer.name }}</a></td>
       <td>{{ peer.address }}</td>
-      <td>{% if peer.enabled %}<span class="badge ok">enabled</span>{% else %}<span class="badge err">disabled</span>{% endif %}</td>
+      <td>{% if peer.enabled %}<span class="badge ok">enabled</span>{% elif peer.over_quota %}<span class="badge err">over quota</span>{% else %}<span class="badge err">disabled</span>{% endif %}</td>
       <td>{% if ps and ps.online %}<span class="badge ok">online</span>{% else %}<span class="badge muted">offline</span>{% endif %}</td>
-      <td class="muted">{{ peer.created_at.strftime("%Y-%m-%d") }}</td>
+      <td class="muted">{{ peer.cum_total | fmt_bytes }}{% if peer.quota_bytes %} / {{ peer.quota_bytes | fmt_bytes }}{% endif %}</td>
+      <td class="muted">{{ peer.note }}</td>
       <td class="actions">
         <form method="post" action="/peers/{{ peer.id }}/toggle" class="inline">
           <button type="submit">{{ "Disable" if peer.enabled else "Enable" }}</button>
@@ -33,7 +38,7 @@
       </td>
     </tr>
   {% else %}
-    <tr><td colspan="6" class="muted">No peers yet. Add one above.</td></tr>
+    <tr><td colspan="7" class="muted">No peers yet. Add one above.</td></tr>
   {% endfor %}
   </tbody>
 </table>