Changeset 2058e5c for server.py

Timestamp:

01/20/26 17:28:10 (6 months ago)

Author:

istevanoska <ilinastevanoska@…>

Branches:

master

Children:

505f39a

Parents:

640ed89

Message:

Working / before login

File:

• server.py (modified) (10 diffs)

server.py

@@ -17 +17 @@
 from openai import OpenAI
 from openai import OpenAI
+from dotenv import load_dotenv
+load_dotenv()
 
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
@@ -25 +27 @@
 
 
-
-
-
 app = Flask(__name__)
 
 DB_FILE = "lan_logs_sysmon.db"
 LOG_DIR = "received_data_sysmon"
+
+import secrets
+
+ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
+ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
+
+
+def init_admin_tables():
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+
+    c.execute('''CREATE TABLE IF NOT EXISTS environments(
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        name TEXT UNIQUE,
+        created_at TEXT
+    )''')
+
+    c.execute('''CREATE TABLE IF NOT EXISTS env_tokens(
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        env_name TEXT,
+        token TEXT UNIQUE,
+        created_at TEXT,
+        expires_at TEXT
+    )''')
+
+    # migration ако табелата веќе постои
+    c.execute("PRAGMA table_info(env_tokens)")
+    cols = [r[1] for r in c.fetchall()]
+    if "expires_at" not in cols:
+        c.execute("ALTER TABLE env_tokens ADD COLUMN expires_at TEXT")
+
+    c.execute('''CREATE TABLE IF NOT EXISTS admin_sessions(
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        token TEXT UNIQUE,
+        created_at TEXT
+    )''')
+
+    # default env
+    c.execute("INSERT OR IGNORE INTO environments(name, created_at) VALUES('default', datetime('now'))")
+
+    conn.commit()
+    conn.close()
+
+def require_admin_session():
+    tok = request.headers.get("X-Admin-Session", "")
+    if not tok:
+        return False
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    c.execute("SELECT 1 FROM admin_sessions WHERE token = ?", (tok,))
+    ok = c.fetchone() is not None
+    conn.close()
+    return ok
+
+def get_env_from_token(req):
+    tok = req.headers.get("X-Env-Token", "")
+    if not tok:
+        return None
+
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    c.execute("""
+        SELECT env_name
+        FROM env_tokens
+        WHERE token = ?
+          AND (expires_at IS NULL OR expires_at > datetime('now'))
+    """, (tok,))
+    row = c.fetchone()
+    conn.close()
+    return row[0] if row else None
 
 
@@ -51 +120 @@
                   last_seen TIMESTAMP,
                   sysmon_available BOOLEAN DEFAULT 0)''')
+    # --- MIGRATION: add env_name column if missing ---
+    c.execute("PRAGMA table_info(computers)")
+    cols = [r[1] for r in c.fetchall()]
+    if "env_name" not in cols:
+        c.execute("ALTER TABLE computers ADD COLUMN env_name TEXT DEFAULT 'default'")
+    # --- END MIGRATION ---
 
     # Табела за историја
@@ -124 +199 @@
         if not data:
             return jsonify({"error": "No data"}), 400
+        env_name = get_env_from_token(request)
+        if not env_name:
+            return jsonify({"error": "Missing or invalid X-Env-Token"}), 401
 
         info = data['info']
@@ -141 +219 @@
             computer_id = result[0]
             c.execute('''UPDATE computers 
-                        SET user = ?, ip = ?, os = ?, last_seen = ?, sysmon_available = ?
+                        SET user = ?, ip = ?, os = ?, last_seen = ?, sysmon_available = ?, env_name = ?
                         WHERE id = ?''',
                       (info['user'], info['ip_address'], info['os'],
                        timestamp.isoformat(), info.get('is_sysmon_available', 0),
-                       computer_id))
+                       env_name, computer_id))
         else:
             # Нов компјутер - додади го
             c.execute('''INSERT INTO computers 
-                        (name, user, ip, os, first_seen, last_seen, sysmon_available) 
-                        VALUES (?, ?, ?, ?, ?, ?, ?)''',
+                        (name, user, ip, os, first_seen, last_seen, sysmon_available, env_name) 
+                        VALUES (?, ?, ?, ?, ?, ?, ?, ?)''',
                       (computer_name, info['user'], info['ip_address'],
                        info['os'], timestamp.isoformat(), timestamp.isoformat(),
-                       info.get('is_sysmon_available', 0)))
+                       info.get('is_sysmon_available', 0), env_name))
+
             computer_id = c.lastrowid
 
@@ -270 +349 @@
 
     # Земaj ги сите компјутери
-    c.execute('''SELECT name, user, ip, os, first_seen, last_seen, sysmon_available 
-                 FROM computers 
+    c.execute('''SELECT name, user, ip, os, first_seen, last_seen, sysmon_available, env_name
+                 FROM computers
                  ORDER BY last_seen DESC''')
 
@@ -331 +410 @@
             'avg_ram': round(avg_ram, 1),
             'status': status,
-            'status_color': status_color
+            'status_color': status_color,
+            'env_name': row[7],
         })
 
@@ -772 +852 @@
                             <i>🌐</i> {{ comp.ip }}
                         </div>
+                        <div class="computer-info">
+                        <i>🏷️</i> Env: {{ comp.env_name }}</div>
+
 
                         <div class="computer-info">
@@ -1324 +1407 @@
     conn = sqlite3.connect(DB_FILE)
     c = conn.cursor()
+    env = request.headers.get("X-Env", "default")
 
     c.execute('''SELECT name, user, ip, os, first_seen, last_seen, sysmon_available 
                  FROM computers 
-                 ORDER BY last_seen DESC''')
+                 WHERE env_name = ?
+                 ORDER BY last_seen DESC''', (env,))
 
     computers = []
@@ -1866 +1951 @@
         'count': len(detailed_processes)
     })
+@app.route("/api/admin/login", methods=["POST"])
+def admin_login():
+    data = request.get_json(force=True) or {}
+    key = (data.get("admin_key") or "").strip()
+    if not ADMIN_KEY or key != ADMIN_KEY:
+        return jsonify({"error": "Unauthorized"}), 401
+
+    session_token = secrets.token_urlsafe(32)
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    c.execute("INSERT INTO admin_sessions(token, created_at) VALUES(?, datetime('now'))", (session_token,))
+    conn.commit()
+    conn.close()
+
+    return jsonify({"ok": True, "session": session_token})
+
+
+@app.route("/api/admin/environments", methods=["GET"])
+def admin_list_envs():
+    if not require_admin_session():
+        return jsonify({"error": "Unauthorized"}), 401
+
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    c.execute("SELECT name FROM environments ORDER BY name")
+    envs = [r[0] for r in c.fetchall()]
+    conn.close()
+
+    return jsonify({"environments": envs})
+
+
+@app.route("/api/admin/environments", methods=["POST"])
+def admin_create_env():
+    if not require_admin_session():
+        return jsonify({"error": "Unauthorized"}), 401
+
+    data = request.get_json(force=True) or {}
+    name = (data.get("name") or "").strip()
+    if not name:
+        return jsonify({"error": "Missing env name"}), 400
+
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+    try:
+        c.execute("INSERT INTO environments(name, created_at) VALUES(?, datetime('now'))", (name,))
+        conn.commit()
+    except Exception as e:
+        conn.close()
+        return jsonify({"error": str(e)}), 400
+    conn.close()
+
+    return jsonify({"ok": True, "name": name})
+
+
+@app.route("/api/admin/tokens", methods=["POST"])
+def admin_generate_token():
+    if not require_admin_session():
+        return jsonify({"error": "Unauthorized"}), 401
+
+    data = request.get_json(force=True) or {}
+    env = (data.get("env") or "").strip()
+    if not env:
+        return jsonify({"error": "Missing env"}), 400
+
+    conn = sqlite3.connect(DB_FILE)
+    c = conn.cursor()
+
+    c.execute("SELECT 1 FROM environments WHERE name = ?", (env,))
+    if not c.fetchone():
+        conn.close()
+        return jsonify({"error": "Environment not found"}), 404
+
+    token = secrets.token_urlsafe(32)
+    c.execute("""
+        INSERT INTO env_tokens(env_name, token, created_at, expires_at)
+        VALUES(?, ?, datetime('now'), datetime('now', '+15 minutes'))
+    """, (env, token))
+
+    conn.commit()
+    conn.close()
+
+    return jsonify({"ok": True, "env": env, "token": token})
+
 if __name__ == '__main__':
     init_database()
+    init_admin_tables()
+
 
     # Стартувај cleanup thread