Refactor process_repo into utils

Author: sylwia-budzynska

src/seclab_taskflows/mcp_servers/codeql_python/mcp_server.py

@@ -21,10 +21,8 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
-
-
-
 from .codeql_sqlite_models import Base, Source
+from .utils import process_repo
 
 MEMORY = Path(os.getenv('CODEQL_SQLITE_DIR', default='/app/my_data'))
 mcp = FastMCP("CodeQL-Python")
@@ -66,7 +64,9 @@ def _resolve_db_path(relative_db_path: str | Path):
     # not windows compatible and probably needs additional hardening
     relative_db_path = str(relative_db_path).strip().lstrip('/')
     relative_db_path = Path(relative_db_path)
-    absolute_path = CODEQL_DBS_BASE_PATH / relative_db_path
+    absolute_path = (CODEQL_DBS_BASE_PATH / relative_db_path).resolve()
+    if not str(absolute_path).startswith(str(CODEQL_DBS_BASE_PATH.resolve())):
+        raise RuntimeError(f"Error: Database path {absolute_path} is outside the base path {CODEQL_DBS_BASE_PATH}")
     if not absolute_path.is_dir():
         _debug_log(f"Database path not found: {absolute_path}")
         raise RuntimeError(f"Error: Database not found at {absolute_path}!")
@@ -148,11 +148,6 @@ def _run_query(query_name: str, database_path: str, language: str, template_valu
     except Exception as e:
         return f"The query {query_name} encountered an error: {e}"
 
-def _get_file_contents(db: str | Path, uri: str):
-    """Retrieve file contents from a CodeQL database"""
-    db = Path(db)
-    return file_from_uri(uri, db)
-
 backend = CodeqlSqliteBackend(MEMORY)
 
 @mcp.tool()
@@ -161,7 +156,7 @@ def remote_sources(owner: str, repo: str,
                    language: str = Field(description="The language used for the CodeQL database.")):
     """List all remote sources and their locations in a CodeQL database, then store the results in a database."""
 
-    repo = f"{owner}/{repo}"
+    repo = process_repo(owner, repo)
     results = _run_query('remote_sources', database_path, language, {})
 
     # Check if results is an error (list of strings) or valid data (list of dicts)
@@ -190,7 +185,7 @@ def fetch_sources(owner: str, repo: str):
     """
     Fetch all sources from the repo
     """
-    repo = f"{owner}/{repo}"
+    repo = process_repo(owner, repo)
     return json.dumps(backend.get_sources(repo))
 
 @mcp.tool()
@@ -202,15 +197,15 @@ def add_source_notes(owner: str, repo: str,
     """
     Add new notes to an existing source. The notes will be appended to any existing notes.
     """
-    repo = f"{owner}/{repo}"
+    repo = process_repo(owner, repo)
     return backend.store_new_source(repo = repo, source_location = source_location, line = line, type = "", notes = notes, update=True)
 
 @mcp.tool()
 def clear_codeql_repo(owner: str, repo: str):
     """
     Clear all data for a given repo from the database
     """
-    repo = f"{owner}/{repo}"
+    repo = process_repo(owner, repo)
     with Session(backend.engine) as session:
         deleted_sources = session.query(Source).filter_by(repo=repo).delete()
         # deleted_apps = session.query(Application).filter_by(repo=repo).delete()

src/seclab_taskflows/mcp_servers/repo_context.py

@@ -21,6 +21,7 @@
 from pathlib import Path
 
 from .repo_context_models import Application, EntryPoint, UserAction, WebEntryPoint, ApplicationIssue, AuditResult, Base
+from .utils import process_repo
 
 MEMORY = Path(os.getenv('REPO_CONTEXT_DIR', default='/app/my_data'))
 
@@ -90,9 +91,9 @@ def __init__(self, memcache_state_dir: str):
         else:
             db_dir = f'sqlite:///{self.memcache_state_dir}/repo_context.db'
         self.engine = create_engine(db_dir, echo=False)
-        Base.metadata.create_all(self.engine, tables=[Application.__table__, EntryPoint.__table__, UserAction.__table__, 
+        Base.metadata.create_all(self.engine, tables=[Application.__table__, EntryPoint.__table__, UserAction.__table__,
                                                       WebEntryPoint.__table__, ApplicationIssue.__table__, AuditResult.__table__])
-    
+
     def store_new_application(self, repo, location, is_app, is_library, notes):
         with Session(self.engine) as session:
             existing = session.query(Application).filter_by(repo = repo, location = location).first()
@@ -107,7 +108,7 @@ def store_new_application(self, repo, location, is_app, is_library, notes):
                 session.add(new_application)
             session.commit()
         return f"Updated or added application for {location} in {repo}."
-    
+
     def store_new_component_issue(self, repo, component_id, issue_type, notes):
         with Session(self.engine) as session:
             existing = session.query(ApplicationIssue).filter_by(repo = repo, component_id = component_id, issue_type = issue_type).first()
@@ -155,7 +156,7 @@ def store_new_entry_point(self, repo, app_id, file, user_input, line, notes, upd
                 session.add(new_entry_point)
             session.commit()
         return f"Updated or added entry point for {file} and {line} in {repo}."
-    
+
     def store_new_web_entry_point(self, repo, entry_point_id, method, path, component, auth, middleware, roles_scopes, notes, update = False):
         with Session(self.engine) as session:
             existing = session.query(WebEntryPoint).filter_by(repo = repo, entry_point_id = entry_point_id).first()
@@ -177,7 +178,7 @@ def store_new_web_entry_point(self, repo, entry_point_id, method, path, componen
                 if update:
                     return f"No web entry point exists at repo {repo} with entry_point_id {entry_point_id}."
                 new_web_entry_point = WebEntryPoint(
-                    repo = repo, 
+                    repo = repo,
                     entry_point_id = entry_point_id,
                     method = method,
                     path = path,
@@ -190,7 +191,7 @@ def store_new_web_entry_point(self, repo, entry_point_id, method, path, componen
                 session.add(new_web_entry_point)
             session.commit()
         return f"Updated or added web entry point for entry_point_id {entry_point_id} in {repo}."
-    
+
     def store_new_user_action(self, repo, app_id, file, line, notes, update = False):
         with Session(self.engine) as session:
             existing = session.query(UserAction).filter_by(repo = repo, file = file, line = line).first()
@@ -203,7 +204,7 @@ def store_new_user_action(self, repo, app_id, file, line, notes, update = False)
                 session.add(new_user_action)
             session.commit()
         return f"Updated or added user action for {file} and {line} in {repo}."
-    
+
     def get_app(self, repo, location):
         with Session(self.engine) as session:
             existing = session.query(Application).filter_by(repo = repo, location = location).first()
@@ -271,7 +272,7 @@ def get_web_entries_for_repo(self, repo):
         with Session(self.engine) as session:
             results = session.query(WebEntryPoint).filter_by(repo = repo).all()
         return [{
-                    'repo' : r.repo, 
+                    'repo' : r.repo,
                     'entry_point_id' : r.entry_point_id,
                     'method' : r.method,
                     'path' : r.path,
@@ -286,7 +287,7 @@ def get_web_entries(self, repo, component_id):
         with Session(self.engine) as session:
             results = session.query(WebEntryPoint).filter_by(repo = repo, component = component_id).all()
         return [{
-                    'repo' : r.repo, 
+                    'repo' : r.repo,
                     'entry_point_id' : r.entry_point_id,
                     'method' : r.method,
                     'path' : r.path,
@@ -313,7 +314,7 @@ def get_user_actions_for_repo(self, repo):
             ).filter(UserAction.app_id == Application.id).all()
             uas = [user_action_to_dict(ua) for app, ua in results]
         return uas
-    
+
     def clear_repo(self, repo):
         with Session(self.engine) as session:
             session.query(Application).filter_by(repo = repo).delete()
@@ -324,7 +325,7 @@ def clear_repo(self, repo):
             session.query(AuditResult).filter_by(repo = repo).delete()
             session.commit()
         return f"Cleared results for repo {repo}"
-    
+
     def clear_repo_issues(self, repo):
         with Session(self.engine) as session:
             session.query(ApplicationIssue).filter_by(repo = repo).delete()
@@ -336,13 +337,10 @@ def clear_repo_issues(self, repo):
 
 backend = RepoContextBackend(MEMORY)
 
-def process_repo(owner, repo):
-    return f"{owner}/{repo}".lower()
-
 @mcp.tool()
-def store_new_component(owner: str, repo: str, location: str = Field(description="The directory of the component"), 
-                  is_app: bool = Field(description="Is this an application", default=None), 
-                  is_library: bool = Field(description="Is this a library", default=None), 
+def store_new_component(owner: str, repo: str, location: str = Field(description="The directory of the component"),
+                  is_app: bool = Field(description="Is this an application", default=None),
+                  is_library: bool = Field(description="Is this a library", default=None),
                   notes: str = Field(description="The notes taken for this component", default="")):
     """
     Stores a new component in the database.
@@ -386,9 +384,9 @@ def store_new_component_issue(owner: str, repo: str, component_id: int,
     return backend.store_new_component_issue(repo, component_id, issue_type, notes)
 
 @mcp.tool()
-def store_new_audit_result(owner: str, repo: str, component_id: int, issue_type: str, issue_id: int, 
-                          has_non_security_error: bool = Field(description="Set to true if there are security issues or logic error but may not be exploitable"), 
-                          has_vulnerability: bool = Field(description="Set to true if a security vulnerability is identified"), 
+def store_new_audit_result(owner: str, repo: str, component_id: int, issue_type: str, issue_id: int,
+                          has_non_security_error: bool = Field(description="Set to true if there are security issues or logic error but may not be exploitable"),
+                          has_vulnerability: bool = Field(description="Set to true if a security vulnerability is identified"),
                           notes: str = Field(description="The notes for the audit of this issue")):
     """
     Stores the audit result for issue with issue_id.
@@ -397,7 +395,7 @@ def store_new_audit_result(owner: str, repo: str, component_id: int, issue_type:
     return backend.store_new_audit_result(repo, component_id, issue_type, issue_id, has_non_security_error, has_vulnerability, notes)
 
 @mcp.tool()
-def store_new_web_entry_point(owner: str, repo: str, 
+def store_new_web_entry_point(owner: str, repo: str,
                                entry_point_id: int = Field(description="The ID of the entry point this web entry point refers to"),
                                location: str = Field(description="The directory of the component where the web entry point belongs to"),
                                method: str = Field(description="HTTP method (GET, POST, etc)", default=""),
@@ -432,7 +430,7 @@ def add_entry_point_notes(owner: str, repo: str,
 @mcp.tool()
 def store_new_user_action(owner: str, repo: str, location: str = Field(description="The directory of the component where the user action belonged to"),
                           file: str = Field(description="The file that contains the user action"),
-                          line: int = Field(description="The file line that contains the user action"), 
+                          line: int = Field(description="The file line that contains the user action"),
                           notes: str = Field(description="New notes for this user action", default = "")):
     """
     Stores a new user action in a component to the database.

src/seclab_taskflows/mcp_servers/utils.py

@@ -0,0 +1,2 @@
+def process_repo(owner, repo):
+    return f"{owner}/{repo}".lower()