Replace mongo with Firestore

This commit is contained in:
Logan Cusano
2025-12-28 12:04:21 -05:00
parent 2e491f8111
commit bfd3b5b2c5
4 changed files with 47 additions and 36 deletions

1
.gitignore vendored
View File

@@ -2,3 +2,4 @@
*.log *.log
*.db *.db
*.conf *.conf
*.json

View File

@@ -1,10 +1,13 @@
import json import json
import os import os
import asyncio import asyncio
from functools import partial
import traceback
from fastapi import FastAPI, HTTPException from fastapi import FastAPI, HTTPException
import paho.mqtt.client as mqtt import paho.mqtt.client as mqtt
from datetime import datetime from datetime import datetime
from motor.motor_asyncio import AsyncIOMotorClient import firebase_admin
from firebase_admin import credentials, firestore
from pydantic import BaseModel from pydantic import BaseModel
from typing import Any, Dict from typing import Any, Dict
@@ -12,13 +15,21 @@ app = FastAPI(title="Radio C2 Brain")
# Configuration # Configuration
MQTT_BROKER = os.getenv("MQTT_BROKER", "mqtt-broker") MQTT_BROKER = os.getenv("MQTT_BROKER", "mqtt-broker")
MONGO_URI = os.getenv("MONGO_URI", "mongodb://admin:securepassword@db:27017/radio_c2?authSource=admin") FIREBASE_CRED_JSON = os.getenv("FIREBASE_CRED_JSON")
FIRESTORE_DB_ID = os.getenv("FIRESTORE_DB_ID", "c2-server")
C2_ID = "central-brain-01" C2_ID = "central-brain-01"
# Database Init # Database Init
mongo_client = AsyncIOMotorClient(MONGO_URI) if FIREBASE_CRED_JSON:
db = mongo_client.get_database() print("Initializing Firebase with provided JSON credentials...")
nodes_col = db.get_collection("nodes") cred = credentials.Certificate(json.loads(FIREBASE_CRED_JSON))
firebase_admin.initialize_app(cred)
else:
print("Initializing Firebase with Application Default Credentials...")
firebase_admin.initialize_app()
print(f"Connecting to Firestore Database: {FIRESTORE_DB_ID}")
db = firestore.client(database_id=FIRESTORE_DB_ID)
# Local cache for quick lookups # Local cache for quick lookups
ACTIVE_NODES_CACHE = {} ACTIVE_NODES_CACHE = {}
@@ -29,6 +40,11 @@ class NodeCommand(BaseModel):
command: str command: str
payload: Dict[str, Any] payload: Dict[str, Any]
# Helper for async execution of blocking firestore calls
async def async_firestore(func, *args, **kwargs):
loop = asyncio.get_running_loop()
return await loop.run_in_executor(None, partial(func, *args, **kwargs))
def on_connect(client, userdata, flags, rc): def on_connect(client, userdata, flags, rc):
print(f"Brain connected to MQTT Broker with result code {rc}") print(f"Brain connected to MQTT Broker with result code {rc}")
client.subscribe("nodes/+/checkin") client.subscribe("nodes/+/checkin")
@@ -50,6 +66,7 @@ async def handle_message(msg):
timestamp = datetime.utcnow() timestamp = datetime.utcnow()
if event_type == "checkin": if event_type == "checkin":
print(f"Processing checkin for {node_id}...")
data = { data = {
"node_id": node_id, "node_id": node_id,
"last_seen": timestamp, "last_seen": timestamp,
@@ -59,20 +76,25 @@ async def handle_message(msg):
"config": payload, "config": payload,
"radio_state": "active" if payload.get("is_listening") else "idle" "radio_state": "active" if payload.get("is_listening") else "idle"
} }
await nodes_col.update_one({"node_id": node_id}, {"$set": data}, upsert=True) doc_ref = db.collection("nodes").document(node_id)
print(f"Writing to Firestore: {doc_ref.path} in DB {FIRESTORE_DB_ID}")
await async_firestore(doc_ref.set, data, merge=True)
print(f"Successfully updated checkin for {node_id}")
ACTIVE_NODES_CACHE[node_id] = data ACTIVE_NODES_CACHE[node_id] = data
elif event_type == "status": elif event_type == "status":
print(f"Processing status update for {node_id}...")
status = payload.get("status") status = payload.get("status")
await nodes_col.update_one( doc_ref = db.collection("nodes").document(node_id)
{"node_id": node_id}, print(f"Writing to Firestore: {doc_ref.path} in DB {FIRESTORE_DB_ID}")
{"$set": {"status": status, "last_seen": timestamp}} await async_firestore(doc_ref.set, {"status": status, "last_seen": timestamp}, merge=True)
) print(f"Successfully updated status for {node_id}")
if node_id in ACTIVE_NODES_CACHE: if node_id in ACTIVE_NODES_CACHE:
ACTIVE_NODES_CACHE[node_id]["status"] = status ACTIVE_NODES_CACHE[node_id]["status"] = status
except Exception as e: except Exception as e:
print(f"Error processing MQTT: {e}") print(f"Error processing MQTT message from {node_id}: {e}")
traceback.print_exc()
# MQTT Setup # MQTT Setup
mqtt_client = mqtt.Client(client_id=C2_ID) mqtt_client = mqtt.Client(client_id=C2_ID)
@@ -88,9 +110,13 @@ async def startup_event():
@app.get("/nodes") @app.get("/nodes")
async def get_nodes(): async def get_nodes():
nodes = await nodes_col.find().to_list(length=100) def get_docs():
# Convert ObjectId to string for JSON serialization return [
for n in nodes: n["_id"] = str(n["_id"]) {**doc.to_dict(), "_id": doc.id}
for doc in db.collection("nodes").stream()
]
nodes = await async_firestore(get_docs)
return nodes return nodes
@app.post("/nodes/{node_id}/command") @app.post("/nodes/{node_id}/command")

View File

@@ -5,30 +5,17 @@ services:
container_name: radio-c2-brain container_name: radio-c2-brain
restart: always restart: always
environment: environment:
- MONGO_URI=mongodb://${MONGO_INITDB_ROOT_USERNAME}:${MONGO_INITDB_ROOT_PASSWORD}@db:27017/radio_c2?authSource=admin - FIREBASE_CRED_JSON=${FIREBASE_CRED_JSON:-}
- FIRESTORE_DB_ID=${FIRESTORE_DB_ID:-c2-server}
- MQTT_BROKER=mqtt-broker - MQTT_BROKER=mqtt-broker
- PORT=8000 - PORT=8000
depends_on: depends_on:
- db
- mqtt-broker - mqtt-broker
ports: ports:
- "8000:8000" - "8000:8000"
networks: networks:
- radio-shared-net - radio-shared-net
# The Memory (MongoDB)
db:
image: mongo:latest
container_name: radio-db
restart: always
environment:
- MONGO_INITDB_ROOT_USERNAME=${MONGO_INITDB_ROOT_USERNAME}
- MONGO_INITDB_ROOT_PASSWORD=${MONGO_INITDB_ROOT_PASSWORD}
volumes:
- mongo_data:/data/db
networks:
- radio-shared-net
# The Post Office (MQTT Broker) # The Post Office (MQTT Broker)
mqtt-broker: mqtt-broker:
image: eclipse-mosquitto:latest image: eclipse-mosquitto:latest
@@ -47,6 +34,3 @@ services:
networks: networks:
radio-shared-net: radio-shared-net:
external: true external: true
volumes:
mongo_data:

View File

@@ -1,6 +1,6 @@
fastapi fastapi
uvicorn[standard] uvicorn[standard]
paho-mqtt paho-mqtt
motor
pydantic pydantic
pydantic-settings pydantic-settings
firebase-admin