Pathc

neuralink/build/service/security_service.py

@@ -8,7 +8,7 @@ import hashlib
 import secrets
 import sqlite3
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional
 
 
 DEFAULT_USERNAME = "default"
@@ -33,6 +33,16 @@ class SecurityService:
             )
             """
         )
+        self._conn.execute(
+            """
+            CREATE TABLE IF NOT EXISTS implants (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                name TEXT UNIQUE NOT NULL,
+                info TEXT NOT NULL,
+                ownername TEXT NOT NULL
+            )
+            """
+        )
         self._conn.commit()
 
     # Utility helpers -----------------------------------------------------
@@ -115,6 +125,41 @@ class SecurityService:
         ).fetchone()
         return row[0] if row else None
 
+    def _user_implants(self, username: str) -> List[str]:
+        rows = self._conn.execute(
+            "SELECT name FROM implants WHERE ownername = ?", (username,)
+        ).fetchall()
+        return [row[0] for row in rows]
+
+    def _implant_belongs_to(self, username: str, name: str) -> bool:
+        row = self._conn.execute(
+            "SELECT 1 FROM implants WHERE ownername = ? AND name = ?",
+            (username, name),
+        ).fetchone()
+        return row is not None
+
+    def _delete_implant(self, name: str) -> bool:
+        cursor = self._conn.execute("DELETE FROM implants WHERE name = ?", (name,))
+        self._conn.commit()
+        return cursor.rowcount > 0
+
+    def _add_implant(self, name: str, info: str, ownername: str) -> bool:
+        try:
+            self._conn.execute(
+                "INSERT INTO implants (name, info, ownername) VALUES (?, ?, ?)",
+                (name, info, ownername),
+            )
+            self._conn.commit()
+        except sqlite3.IntegrityError:
+            return False
+        return True
+
+    def _get_implant_info(self, name: str) -> Optional[str]:
+        row = self._conn.execute(
+            "SELECT info FROM implants WHERE name = ?", (name,)
+        ).fetchone()
+        return row[0] if row else None
+
     # User flows ----------------------------------------------------------
     def register_user(self) -> None:
         raw_username = input("\nEnter username: ").strip()
@@ -196,6 +241,82 @@ class SecurityService:
         else:
             print(f"Security code: {code}")
 
+    def add_implant(self) -> None:
+        if self.username == DEFAULT_USERNAME:
+            self.print_system_data("You need to log in first.")
+            return
+
+        name = input("\nEnter implant name: ").strip()
+        info = input("\nEnter implant info: ").strip()
+        if not name or not info:
+            self.print_system_data("Invalid option selected.")
+            return
+
+        if self._add_implant(name, info, self.username):
+            self.print_system_data("Implant added successfully.")
+        else:
+            self.print_system_data("Unexpected error. Please try later.")
+
+    def delete_implant(self) -> None:
+        if self.username == DEFAULT_USERNAME:
+            self.print_system_data("You need to log in first.")
+            return
+
+        self.print_system_data("Getting list of implants. Please wait...")
+        implants = self._user_implants(self.username)
+        if not implants:
+            self.print_system_data("No implants found for this user.")
+            return
+
+        for idx, implant in enumerate(implants, start=1):
+            print(f"{idx}. {implant}")
+
+        name = input("\nEnter implant name: ").strip()
+        if not name:
+            self.print_system_data("Invalid option selected.")
+            return
+
+        if not self._implant_belongs_to(self.username, name):
+            self.print_system_data("Invalid option selected.")
+            return
+
+        self.print_system_data("Deleting implant...")
+        if self._delete_implant(name):
+            self.print_system_data("Implant deleted successfully.")
+        else:
+            self.print_system_data("Unexpected error. Please try later.")
+
+    def show_implant_info(self) -> None:
+        if self.username == DEFAULT_USERNAME:
+            self.print_system_data("You need to log in first.")
+            return
+
+        self.print_system_data("Getting list of implants. Please wait...")
+        implants = self._user_implants(self.username)
+        if not implants:
+            self.print_system_data("No implants found for this user.")
+            return
+
+        for idx, implant in enumerate(implants, start=1):
+            print(f"{idx}. {implant}")
+
+        name = input("\nEnter implant name: ").strip()
+        if not name:
+            self.print_system_data("Invalid option selected.")
+            return
+
+        if not self._implant_belongs_to(self.username, name):
+            self.print_system_data("Invalid option selected.")
+            return
+
+        self.print_system_data("Getting implant info...")
+        info = self._get_implant_info(name)
+        if info is None:
+            self.print_system_data("Unexpected error. Please try later.")
+            return
+
+        print(f"Implant info: {info}")
+
     # Menus ---------------------------------------------------------------
     def settings_menu(self) -> None:
         while self.username != DEFAULT_USERNAME:
@@ -216,17 +337,39 @@ class SecurityService:
     def app_menu(self) -> None:
         while self.username != DEFAULT_USERNAME:
             self.print_system_data("--- Application menu ---")
-            self.print_system_data("1. Settings menu")
-            self.print_system_data("2. Log out")
-            choice = input("Choose an option (1-2): ").strip()
+            self.print_system_data("1. Implant menu")
+            self.print_system_data("2. Settings menu")
+            self.print_system_data("3. Log out")
+            choice = input("Choose an option (1-3): ").strip()
             if choice == "1":
-                self.settings_menu()
+                self.implants_menu()
             elif choice == "2":
+                self.settings_menu()
+            elif choice == "3":
                 self.username = DEFAULT_USERNAME
                 self.print_system_data("Logged out.")
             else:
                 self.print_system_data("Invalid option selected.")
 
+    def implants_menu(self) -> None:
+        while self.username != DEFAULT_USERNAME:
+            self.print_system_data("--- Implants menu ---")
+            self.print_system_data("1. Add new implant")
+            self.print_system_data("2. Delete implant")
+            self.print_system_data("3. Show info about implant")
+            self.print_system_data("4. Return to previous menu")
+            choice = input("Choose an option (1-4): ").strip()
+            if choice == "1":
+                self.add_implant()
+            elif choice == "2":
+                self.delete_implant()
+            elif choice == "3":
+                self.show_implant_info()
+            elif choice == "4":
+                return
+            else:
+                self.print_system_data("Invalid option selected.")
+
     def startup_menu(self) -> None:
         while True:
             if self.username != DEFAULT_USERNAME: