Replace mongo with Firestore
This commit is contained in:
3
.gitignore
vendored
3
.gitignore
vendored
@@ -1,4 +1,5 @@
|
|||||||
.env
|
.env
|
||||||
*.log
|
*.log
|
||||||
*.db
|
*.db
|
||||||
*.conf
|
*.conf
|
||||||
|
*.json
|
||||||
@@ -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")
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -46,7 +33,4 @@ services:
|
|||||||
|
|
||||||
networks:
|
networks:
|
||||||
radio-shared-net:
|
radio-shared-net:
|
||||||
external: true
|
external: true
|
||||||
|
|
||||||
volumes:
|
|
||||||
mongo_data:
|
|
||||||
@@ -1,6 +1,6 @@
|
|||||||
fastapi
|
fastapi
|
||||||
uvicorn[standard]
|
uvicorn[standard]
|
||||||
paho-mqtt
|
paho-mqtt
|
||||||
motor
|
|
||||||
pydantic
|
pydantic
|
||||||
pydantic-settings
|
pydantic-settings
|
||||||
|
firebase-admin
|
||||||
Reference in New Issue
Block a user