334 lines
13 KiB
Python
334 lines
13 KiB
Python
|
import json
|
||
|
|
import os
|
||
|
|
import secrets
|
||
|
|
import sqlite3
|
||
|
|
import time
|
||
|
|
import urllib.request
|
||
|
|
import urllib.error
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from functools import wraps
|
||
|
|
from pathlib import Path
|
||
|
|
from threading import Thread
|
||
|
|
|
||
|
|
import hmac, hashlib, base64
|
||
|
|
|
||
|
|
from flask import Flask, request, jsonify
|
||
|
|
from flask_cors import CORS
|
||
|
|
from nacl.signing import VerifyKey
|
||
|
|
from nacl.exceptions import BadSignatureError
|
||
|
|
|
||
|
|
app = Flask(__name__)
|
||
|
|
CORS(app)
|
||
|
|
|
||
|
|
DB_PATH = os.environ.get("MONTANA_DB_PATH", "/var/lib/montana-messenger/db.sqlite3")
|
||
|
|
PEERS_PATH = os.environ.get("MONTANA_PEERS_PATH", "/etc/montana-site/peers.json")
|
||
|
|
SHARED_SECRET = os.environ.get("MONTANA_SHARED_SECRET", "change-me-in-prod")
|
||
|
|
NODE_NAME = os.environ.get("MONTANA_NODE_NAME", "unknown")
|
||
|
|
|
||
|
|
CHALLENGE_TTL_MS = 5 * 60 * 1000
|
||
|
|
SESSION_TTL_MS = 30 * 24 * 60 * 60 * 1000
|
||
|
|
|
||
|
|
|
||
|
|
def now_ms() -> int:
|
||
|
|
return int(time.time() * 1000)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def db():
|
||
|
|
Path(DB_PATH).parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
conn = sqlite3.connect(DB_PATH, isolation_level=None, timeout=10)
|
||
|
|
conn.row_factory = sqlite3.Row
|
||
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
||
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
||
|
|
try:
|
||
|
|
yield conn
|
||
|
|
finally:
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
|
||
|
|
def init_schema():
|
||
|
|
schema = (Path(__file__).parent / "schema.sql").read_text()
|
||
|
|
with db() as conn:
|
||
|
|
conn.executescript(schema)
|
||
|
|
|
||
|
|
|
||
|
|
def load_peers():
|
||
|
|
try:
|
||
|
|
return json.loads(Path(PEERS_PATH).read_text()).get("peers", [])
|
||
|
|
except Exception:
|
||
|
|
return []
|
||
|
|
|
||
|
|
|
||
|
|
# ─── JWT-like signed challenge (stateless, любой узел может verify) ───
|
||
|
|
def b64u_encode(b: bytes) -> str:
|
||
|
|
return base64.urlsafe_b64encode(b).rstrip(b"=").decode("ascii")
|
||
|
|
|
||
|
|
|
||
|
|
def b64u_decode(s: str) -> bytes:
|
||
|
|
pad = "=" * (-len(s) % 4)
|
||
|
|
return base64.urlsafe_b64decode(s + pad)
|
||
|
|
|
||
|
|
|
||
|
|
def make_challenge_token(account_id: str, challenge: bytes, expires_at: int) -> str:
|
||
|
|
payload = json.dumps({"a": account_id, "c": challenge.hex(), "e": expires_at}, separators=(",", ":")).encode()
|
||
|
|
body = b64u_encode(payload)
|
||
|
|
sig = hmac.new(SHARED_SECRET.encode(), body.encode(), hashlib.sha256).digest()
|
||
|
|
return body + "." + b64u_encode(sig)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_challenge_token(token: str):
|
||
|
|
try:
|
||
|
|
body, sig = token.split(".", 1)
|
||
|
|
expected = b64u_encode(hmac.new(SHARED_SECRET.encode(), body.encode(), hashlib.sha256).digest())
|
||
|
|
if not hmac.compare_digest(sig, expected):
|
||
|
|
return None
|
||
|
|
data = json.loads(b64u_decode(body))
|
||
|
|
if data["e"] < now_ms():
|
||
|
|
return None
|
||
|
|
return data["a"], bytes.fromhex(data["c"]), data["e"]
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
# ─── Replicate events (gossip) ───
|
||
|
|
def replicate_async(event: dict):
|
||
|
|
"""Бросает событие всем peer-узлам в фоновом потоке."""
|
||
|
|
Thread(target=_replicate_sync, args=(event,), daemon=True).start()
|
||
|
|
|
||
|
|
|
||
|
|
def _replicate_sync(event: dict):
|
||
|
|
peers = load_peers()
|
||
|
|
body = json.dumps(event).encode()
|
||
|
|
sig = hmac.new(SHARED_SECRET.encode(), body, hashlib.sha256).hexdigest()
|
||
|
|
for peer in peers:
|
||
|
|
url = peer["url"] + "/api/internal/replicate"
|
||
|
|
try:
|
||
|
|
req = urllib.request.Request(
|
||
|
|
url,
|
||
|
|
data=body,
|
||
|
|
headers={"Content-Type": "application/json", "X-Montana-Sig": sig, "X-Montana-Origin": NODE_NAME},
|
||
|
|
method="POST",
|
||
|
|
)
|
||
|
|
urllib.request.urlopen(req, timeout=5)
|
||
|
|
except Exception:
|
||
|
|
pass # fire-and-forget; узел догонит при следующем write
|
||
|
|
|
||
|
|
|
||
|
|
def apply_event(event: dict):
|
||
|
|
"""Применяет реплицируемое событие к локальной БД."""
|
||
|
|
kind = event.get("kind")
|
||
|
|
if kind == "account":
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO accounts (account_id, encryption_pubkey, created_at_ms, last_seen_at_ms) VALUES (?, ?, ?, ?) "
|
||
|
|
"ON CONFLICT(account_id) DO UPDATE SET last_seen_at_ms=MAX(accounts.last_seen_at_ms, excluded.last_seen_at_ms), "
|
||
|
|
"encryption_pubkey=COALESCE(excluded.encryption_pubkey, accounts.encryption_pubkey)",
|
||
|
|
(event["account_id"], event.get("encryption_pubkey"), event.get("created_at_ms", now_ms()), event.get("last_seen_at_ms", now_ms())),
|
||
|
|
)
|
||
|
|
elif kind == "session":
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO sessions (token, account_id, expires_at_ms, created_at_ms) VALUES (?, ?, ?, ?) "
|
||
|
|
"ON CONFLICT(token) DO NOTHING",
|
||
|
|
(event["token"], event["account_id"], event["expires_at_ms"], event["created_at_ms"]),
|
||
|
|
)
|
||
|
|
elif kind == "message":
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO accounts (account_id, created_at_ms) VALUES (?, ?) ON CONFLICT(account_id) DO NOTHING",
|
||
|
|
(event["recipient"], event.get("created_at_ms", now_ms())),
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO messages (msg_id, sender, recipient, ciphertext, nonce, created_at_ms) VALUES (?, ?, ?, ?, ?, ?) "
|
||
|
|
"ON CONFLICT(msg_id) DO NOTHING",
|
||
|
|
(event["msg_id"], event["sender"], event["recipient"], event["ciphertext"], event["nonce"], event["created_at_ms"]),
|
||
|
|
)
|
||
|
|
elif kind == "read":
|
||
|
|
with db() as conn:
|
||
|
|
placeholders = ",".join("?" * len(event["msg_ids"]))
|
||
|
|
conn.execute(
|
||
|
|
f"UPDATE messages SET read_at_ms=? WHERE msg_id IN ({placeholders}) AND recipient=?",
|
||
|
|
(event["read_at_ms"], *event["msg_ids"], event["account_id"]),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# ─── auth_required decorator ───
|
||
|
|
def auth_required(fn):
|
||
|
|
@wraps(fn)
|
||
|
|
def wrapper(*a, **kw):
|
||
|
|
h = request.headers.get("Authorization", "")
|
||
|
|
if not h.startswith("Bearer "):
|
||
|
|
return jsonify({"error": "no_auth"}), 401
|
||
|
|
token = h[7:]
|
||
|
|
with db() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"SELECT account_id FROM sessions WHERE token=? AND expires_at_ms > ?",
|
||
|
|
(token, now_ms()),
|
||
|
|
).fetchone()
|
||
|
|
if not row:
|
||
|
|
return jsonify({"error": "invalid_token"}), 401
|
||
|
|
request.account_id = row["account_id"]
|
||
|
|
return fn(*a, **kw)
|
||
|
|
|
||
|
|
return wrapper
|
||
|
|
|
||
|
|
|
||
|
|
# ─── Public API ───
|
||
|
|
@app.post("/api/auth/challenge")
|
||
|
|
def auth_challenge():
|
||
|
|
data = request.get_json(force=True) or {}
|
||
|
|
account_id = (data.get("accountId") or "").lower()
|
||
|
|
if len(account_id) != 64:
|
||
|
|
return jsonify({"error": "bad_account_id"}), 400
|
||
|
|
challenge = secrets.token_bytes(32)
|
||
|
|
expires_at = now_ms() + CHALLENGE_TTL_MS
|
||
|
|
token = make_challenge_token(account_id, challenge, expires_at)
|
||
|
|
return jsonify({"challengeId": token, "challenge": challenge.hex(), "expiresAt": expires_at, "node": NODE_NAME})
|
||
|
|
|
||
|
|
|
||
|
|
@app.post("/api/auth/verify")
|
||
|
|
def auth_verify():
|
||
|
|
data = request.get_json(force=True) or {}
|
||
|
|
token = data.get("challengeId", "")
|
||
|
|
sig_hex = data.get("signature", "")
|
||
|
|
account_id_claimed = (data.get("accountId") or "").lower()
|
||
|
|
enc_pub = (data.get("encryptionPubkey") or "").lower() or None
|
||
|
|
|
||
|
|
parsed = parse_challenge_token(token)
|
||
|
|
if not parsed:
|
||
|
|
return jsonify({"error": "challenge_expired_or_invalid"}), 400
|
||
|
|
aid_in_token, challenge_bytes, expires = parsed
|
||
|
|
if aid_in_token != account_id_claimed:
|
||
|
|
return jsonify({"error": "account_mismatch"}), 400
|
||
|
|
|
||
|
|
try:
|
||
|
|
VerifyKey(bytes.fromhex(account_id_claimed)).verify(challenge_bytes, bytes.fromhex(sig_hex))
|
||
|
|
except (BadSignatureError, ValueError):
|
||
|
|
return jsonify({"error": "bad_signature"}), 401
|
||
|
|
|
||
|
|
session_token = secrets.token_urlsafe(32)
|
||
|
|
session_expires = now_ms() + SESSION_TTL_MS
|
||
|
|
created = now_ms()
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO accounts (account_id, encryption_pubkey, created_at_ms, last_seen_at_ms) VALUES (?, ?, ?, ?) "
|
||
|
|
"ON CONFLICT(account_id) DO UPDATE SET last_seen_at_ms=excluded.last_seen_at_ms, "
|
||
|
|
"encryption_pubkey=COALESCE(excluded.encryption_pubkey, accounts.encryption_pubkey)",
|
||
|
|
(account_id_claimed, enc_pub, created, created),
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO sessions (token, account_id, expires_at_ms, created_at_ms) VALUES (?, ?, ?, ?)",
|
||
|
|
(session_token, account_id_claimed, session_expires, created),
|
||
|
|
)
|
||
|
|
replicate_async({"kind": "account", "account_id": account_id_claimed, "encryption_pubkey": enc_pub, "created_at_ms": created, "last_seen_at_ms": created})
|
||
|
|
replicate_async({"kind": "session", "token": session_token, "account_id": account_id_claimed, "expires_at_ms": session_expires, "created_at_ms": created})
|
||
|
|
return jsonify({"sessionToken": session_token, "accountId": account_id_claimed, "expiresAt": session_expires, "node": NODE_NAME})
|
||
|
|
|
||
|
|
|
||
|
|
@app.post("/api/msg/send")
|
||
|
|
@auth_required
|
||
|
|
def msg_send():
|
||
|
|
data = request.get_json(force=True) or {}
|
||
|
|
recipient = (data.get("recipient") or "").lower()
|
||
|
|
ciphertext = data.get("ciphertext", "")
|
||
|
|
nonce = data.get("nonce", "")
|
||
|
|
if len(recipient) != 64 or not ciphertext or not nonce:
|
||
|
|
return jsonify({"error": "bad_input"}), 400
|
||
|
|
msg_id = secrets.token_hex(16)
|
||
|
|
created = now_ms()
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute("INSERT INTO accounts (account_id, created_at_ms) VALUES (?, ?) ON CONFLICT(account_id) DO NOTHING", (recipient, created))
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO messages (msg_id, sender, recipient, ciphertext, nonce, created_at_ms) VALUES (?, ?, ?, ?, ?, ?)",
|
||
|
|
(msg_id, request.account_id, recipient, ciphertext, nonce, created),
|
||
|
|
)
|
||
|
|
replicate_async({"kind": "message", "msg_id": msg_id, "sender": request.account_id, "recipient": recipient, "ciphertext": ciphertext, "nonce": nonce, "created_at_ms": created})
|
||
|
|
return jsonify({"msgId": msg_id, "createdAt": created, "node": NODE_NAME})
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/api/msg/inbox")
|
||
|
|
@auth_required
|
||
|
|
def msg_inbox():
|
||
|
|
since = int(request.args.get("since", "0"))
|
||
|
|
with db() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"SELECT msg_id, sender, recipient, ciphertext, nonce, created_at_ms, read_at_ms "
|
||
|
|
"FROM messages WHERE (recipient=? OR sender=?) AND created_at_ms > ? ORDER BY created_at_ms ASC LIMIT 1000",
|
||
|
|
(request.account_id, request.account_id, since),
|
||
|
|
).fetchall()
|
||
|
|
return jsonify({"messages": [{"msgId": r["msg_id"], "sender": r["sender"], "recipient": r["recipient"], "ciphertext": r["ciphertext"], "nonce": r["nonce"], "createdAt": r["created_at_ms"], "readAt": r["read_at_ms"]} for r in rows]})
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/api/msg/contacts")
|
||
|
|
@auth_required
|
||
|
|
def msg_contacts():
|
||
|
|
with db() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"SELECT CASE WHEN sender=? THEN recipient ELSE sender END AS aid, MAX(created_at_ms) AS last "
|
||
|
|
"FROM messages WHERE sender=? OR recipient=? GROUP BY 1 ORDER BY 2 DESC LIMIT 200",
|
||
|
|
(request.account_id, request.account_id, request.account_id),
|
||
|
|
).fetchall()
|
||
|
|
return jsonify({"contacts": [{"accountId": r["aid"], "lastMsgAt": r["last"]} for r in rows]})
|
||
|
|
|
||
|
|
|
||
|
|
@app.post("/api/msg/read")
|
||
|
|
@auth_required
|
||
|
|
def msg_read():
|
||
|
|
data = request.get_json(force=True) or {}
|
||
|
|
msg_ids = data.get("msgIds") or []
|
||
|
|
if not isinstance(msg_ids, list) or not msg_ids:
|
||
|
|
return jsonify({"error": "bad_input"}), 400
|
||
|
|
ts = now_ms()
|
||
|
|
placeholders = ",".join("?" * len(msg_ids))
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute(f"UPDATE messages SET read_at_ms=? WHERE msg_id IN ({placeholders}) AND recipient=?", (ts, *msg_ids, request.account_id))
|
||
|
|
replicate_async({"kind": "read", "msg_ids": msg_ids, "account_id": request.account_id, "read_at_ms": ts})
|
||
|
|
return jsonify({"ok": True})
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/api/account/<account_id>")
|
||
|
|
@auth_required
|
||
|
|
def account_info(account_id):
|
||
|
|
account_id = account_id.lower()
|
||
|
|
if len(account_id) != 64:
|
||
|
|
return jsonify({"error": "bad_account_id"}), 400
|
||
|
|
with db() as conn:
|
||
|
|
row = conn.execute("SELECT account_id, encryption_pubkey, last_seen_at_ms FROM accounts WHERE account_id=?", (account_id,)).fetchone()
|
||
|
|
if not row:
|
||
|
|
return jsonify({"accountId": account_id, "encryptionPubkey": None, "lastSeenAt": None})
|
||
|
|
return jsonify({"accountId": row["account_id"], "encryptionPubkey": row["encryption_pubkey"], "lastSeenAt": row["last_seen_at_ms"]})
|
||
|
|
|
||
|
|
|
||
|
|
# ─── Internal: replication endpoint ───
|
||
|
|
@app.post("/api/internal/replicate")
|
||
|
|
def replicate_in():
|
||
|
|
raw = request.get_data()
|
||
|
|
sig = request.headers.get("X-Montana-Sig", "")
|
||
|
|
expected = hmac.new(SHARED_SECRET.encode(), raw, hashlib.sha256).hexdigest()
|
||
|
|
if not hmac.compare_digest(sig, expected):
|
||
|
|
return jsonify({"error": "bad_sig"}), 401
|
||
|
|
try:
|
||
|
|
event = json.loads(raw)
|
||
|
|
apply_event(event)
|
||
|
|
except Exception as e:
|
||
|
|
return jsonify({"error": str(e)}), 400
|
||
|
|
return jsonify({"ok": True, "node": NODE_NAME})
|
||
|
|
|
||
|
|
|
||
|
|
@app.get("/api/health")
|
||
|
|
def health():
|
||
|
|
try:
|
||
|
|
with db() as conn:
|
||
|
|
conn.execute("SELECT 1").fetchone()
|
||
|
|
peers = [p["name"] for p in load_peers()]
|
||
|
|
return jsonify({"ok": True, "service": "montana-messenger-relay", "node": NODE_NAME, "peers": peers, "ts": now_ms()})
|
||
|
|
except Exception as e:
|
||
|
|
return jsonify({"ok": False, "error": str(e)}), 500
|
||
|
|
|
||
|
|
|
||
|
|
init_schema()
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
app.run(host="127.0.0.1", port=5010, debug=False)
|