Use local resources in CodeQL for Python MCP

Author: sylwia-budzynska

src/seclab_taskflows/mcp_servers/codeql_python/codeql_sqlite_models.py

@@ -15,6 +15,7 @@ class Source(Base):
     id: Mapped[int] = mapped_column(primary_key=True)
     repo: Mapped[str]
     source_location: Mapped[str]
+    line: Mapped[int]
     type: Mapped[str]
     notes: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
 

src/seclab_taskflows/mcp_servers/codeql_python/mcp_server.py

@@ -45,6 +45,7 @@ def source_to_dict(result):
         "source_id": result.id,
         "repo": result.repo,
         "source_location": result.source_location,
+        "line": result.line,
         "type": result.type,
         "notes": result.notes
     }
@@ -84,17 +85,17 @@ def __init__(self, memcache_state_dir: str):
         Base.metadata.create_all(self.engine, tables=[Source.__table__])
 
 
-    def store_new_source(self, repo, source_location, type, notes, update = False):
+    def store_new_source(self, repo, source_location, line, type, notes, update = False):
         with Session(self.engine) as session:
-            existing = session.query(Source).filter_by(repo = repo, source_location = source_location).first()
+            existing = session.query(Source).filter_by(repo = repo, source_location = source_location, line = line).first()
             if existing:
                 existing.notes = (existing.notes or "") + notes
                 session.commit()
-                return f"Updated notes for source at {source_location} in {repo}."
+                return f"Updated notes for source at {source_location}, line {line} in {repo}."
             else:
                 if update:
-                    return f"No source exists at repo {repo}, location {source_location}"
-                new_source = Source(repo = repo,  source_location = source_location, type = type, notes = notes)
+                    return f"No source exists at repo {repo}, location {source_location}, line {line} to update."
+                new_source = Source(repo = repo,  source_location = source_location, line = line, type = type, notes = notes)
                 session.add(new_source)
                 session.commit()
                 return f"Added new source for {source_location} in {repo}."
@@ -174,7 +175,8 @@ def remote_sources(owner: str, repo: str,
             repo=repo,
             source_location=result.get('location', ''),
             type=result.get('source', ''),
-            notes='', #result.get('description', ''),
+            line=int(result.get('line', '0')),
+            notes=None, #result.get('description', ''),
             update=False
         )
         stored_count += 1
@@ -191,18 +193,15 @@ def fetch_sources(owner: str, repo: str):
 
 @mcp.tool()
 def add_source_notes(owner: str, repo: str,
-                     database_path: str = Field(description="The CodeQL database path."),
-                     source_location: str = Field(description="The path to the file and column info that contains the source"),
+                    #  database_path: str = Field(description="The CodeQL database path."),
+                     source_location: str = Field(description="The path to the file"),
+                     line: int = Field(description="The line number of the source"),
                      notes: str = Field(description="The notes to append to this source", default="")):
     """
     Add new notes to an existing source. The notes will be appended to any existing notes.
     """
     repo = f"{owner}/{repo}"
-    try:
-        database_path = _resolve_db_path(database_path)
-    except RuntimeError:
-        return f"The database path for {database_path} could not be resolved"
-    return backend.store_new_source(repo, source_location, "", notes, update=True)
+    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):
@@ -216,45 +215,5 @@ def clear_codeql_repo(owner: str, repo: str):
         session.commit()
     return f"Cleared {deleted_sources} sources from repo {repo}."
 
-@mcp.tool()
-def get_file_contents(
-        file_uri: str = Field(description="The file URI to get contents for. The URI scheme is defined as `file://path` and `file://path:region`. Examples of file URI: `file:///path/to/file:1:2:3:4`, `file:///path/to/file`. File URIs optionally contain a region definition that looks like `start_line:start_column:end_line:end_column` which will limit the contents returned to the specified region, for example `file:///path/to/file:1:2:3:4` indicates a file region of `1:2:3:4` which would return the content of the file starting at line 1, column 1 and ending at line 3 column 4. Line and column indices are 1-based, meaning line and column values start at 1. If the region is omitted the full contents of the file will be returned, for example `file:///path/to/file` returns the full contents of `/path/to/file`."),
-        database_path: str = Field(description="The path to the CodeQL database.")):
-    """Get the contents of a file URI from a CodeQL database path."""
-
-    database_path = _resolve_db_path(database_path)
-    try:
-        # fix up any incorrectly formatted relative path uri
-        if not file_uri.startswith('file:///'):
-            if file_uri.startswith('file://'):
-                file_uri = file_uri[len('file://'):]
-            file_uri = 'file:///' + file_uri.lstrip('/')
-        results = _get_file_contents(database_path, file_uri)
-    except Exception as e:
-        results = f"Error: could not retrieve {file_uri}: {e}"
-    return results
-
-@mcp.tool()
-def list_source_files(database_path: str = Field(description="The path to the CodeQL database."),
-                      regex_filter: str = Field(description="Optional Regex filter.", default = r'[\s\S]+')):
-    """List the available source files in a CodeQL database using their file:// URI"""
-    database_path = _resolve_db_path(database_path)
-    results = list_src_files(database_path, as_uri=True)
-    return json.dumps([{'uri': item} for item in results if re.search(regex_filter, item)], indent=2)
-
-@mcp.tool()
-def search_in_source_code(database_path: str = Field(description="The path to the CodeQL database."),
-                          search_term: str = Field(description="The term to search in the source code")):
-    """
-    Search for a string in the source code. Returns the line number and file.
-    """
-    resolved_database_path = _resolve_db_path(database_path)
-    results = search_in_src_archive(resolved_database_path, search_term)
-    out = []
-    if isinstance(results, dict):
-        for k,v in results.items():
-            out.append({"database" : database_path, "path" : k, "lines" : v})
-    return json.dumps(out, indent = 2)
-
 if __name__ == "__main__":
     mcp.run(show_banner=False, transport="http", host="127.0.0.1", port=9998)

src/seclab_taskflows/mcp_servers/codeql_python/queries/mcp-python/remote_sources.ql

@@ -7,11 +7,15 @@
 import python
 import semmle.python.dataflow.new.RemoteFlowSources
 
-string normalizeLocation(Location l) {
-    result = "file://" + "/" + l.getFile().getRelativePath() + ":" + l.getStartLine().toString() + ":" + l.getStartColumn().toString()
-    + ":" + l.getEndLine().toString() + ":" + l.getEndColumn().toString()
-}
+// string normalizeLocation(Location l) {
+//     result = l.getFile().getRelativePath() + ":" + l.getStartLine().toString() + ":" + l.getStartColumn().toString()
+//     + ":" + l.getEndLine().toString() + ":" + l.getEndColumn().toString()
+// }
 
 from RemoteFlowSource source
 select
-  "Remote source {0} is defined at {1}", "source,location", source.getSourceType(), normalizeLocation(source.getLocation())
+  "Remote source {0} is defined at {1} line {2}",
+  "source,location,line",
+  source.getSourceType(),
+  source.getLocation().getFile().getRelativePath(),
+  source.getLocation().getStartLine().toString()

…lab_taskflows/personalities/auditer.yaml …lab_taskflows/personalities/auditor.yamlsrc/seclab_taskflows/personalities/auditer.yaml renamed to src/seclab_taskflows/personalities/auditor.yaml src/seclab_taskflows/personalities/auditer.yaml renamed to src/seclab_taskflows/personalities/auditor.yaml

@@ -9,7 +9,6 @@ model_config: seclab_taskflows.configs.model_config
 
 globals:
   repo:
-    apache/allura
 # Taskflow to analyze the existing information
 taskflow:
   - task:
@@ -31,11 +30,9 @@ taskflow:
         - seclab_taskflow_agent.personalities.assistant
       user_prompt: |
         For the repo {{ GLOBALS_repo }} fetch the Python CodeQL database and find all remote flow sources using CodeQL.
-        Store the value for CodeQL's 'relative_database_path' in 'codeql_relative_database_path' memory key.
       toolboxes:
         - seclab_taskflows.toolboxes.gh_code_scanning
         - seclab_taskflows.toolboxes.codeql_python
-        - seclab_taskflow_agent.toolboxes.memcache
   - task:
       must_complete: true
       exclude_from_context: true
@@ -50,25 +47,21 @@ taskflow:
       model: code_analysis
       must_complete: false
       repeat_prompt: true
-      async: true
-      async_limit: 5
       max_steps: 100
       name: source analysis
       description: Identify actions that untrusted users are allowed perform the source.
       agents:
-        - seclab_taskflows.personalities.auditer
+        - seclab_taskflows.personalities.auditor
       user_prompt: |
-        Retrieve the contents of the `codeql_relative_database_path` memory key, which represents the relative path to the CodeQL database for the repository {{ GLOBALS_repo }}.
-        Using the CodeQL database located at that path, analyze the following source:
-        The source is a {{ RESULT_type }} in {{ RESULT_repo }} in the location {{ RESULT_source_location }}.
-        Analyze what the source endpoint is for and how it is used.
-        Search for relevant code associated with each source.
+        Fetch the zipball of the repository {{ GLOBALS_repo }} and use it to analyze the source.
+        The source is a {{ RESULT_type }} in {{ RESULT_repo }} in the location {{ RESULT_source_location }} on line {{ RESULT_line }}.
+        If the source is in a folder relating to tests or demo code, skip the analysis and update the source entry in the codeql_sqlite database indicating it is not relevant.
+        Analyze what the source endpoint is used for.
         If it is a web endpoint, identify the routing path that reaches this source, HTTP method,
         any middlewares used, which roles are allowed to call it.
         Note which kind of authentication is required for that endpoint.
         It is possible that the source does not have require any authentication.
         If authorization is required, note the details.
-        Analyze the code and identify if this source could lead to a security vulnerability.
 
         Update the source entry in the codeql_sqlite database with your findings.
         ## IMPORTANT: General Guidance that ALWAYS applies
@@ -82,7 +75,8 @@ taskflow:
         the error you encountered.
       toolboxes:
         - seclab_taskflows.toolboxes.codeql_python
-        - seclab_taskflow_agent.toolboxes.memcache
+        - seclab_taskflows.toolboxes.local_gh_resources
+        - seclab_taskflows.toolboxes.local_file_viewer
   - task:
       must_complete: true
       agents:

src/seclab_taskflows/taskflows/audit/remote_sources_local.yaml

@@ -46,34 +46,3 @@ server_prompt: |
   If you are unable to determine the appropriate programming language acronym,
   halt your task and ask the user to clarify which programming language the
   CodeQL database in question was created for.
-
-
-  ### CodeQL Database File URI
-
-  The CodeQL database file URI scheme is defined as `file://path` and
-  `file://path:region`.
-
-  Examples of CodeQL database file URIs:
-
-    - `file:///path/to/file:1:2:3:4`
-    - `file:///path/to/file`
-
-  File URIs optionally contain a region definition that looks like
-  `start_line:start_column:end_line:end_column` which will limit the contents
-  returned to the specified region. For example `file:///path/to/file:1:2:3:4`
-  indicates a file region of `1:2:3:4` which would return the content of the
-  file starting at line 1, column 2 and ending at line 3, column 4. These line
-  and column indices are 1-based, meaning line and column values start at 1.
-
-  If the region is ommitted the full contents of the file will be returned,
-  for example `file:///path/to/file` returns the full contents of
-  `/path/to/file`.
-
-  If you want to fetch a specific region by their line numbers only, you can set
-  the `start_column` and `end_column` values of a region to `0`. For example to
-  retrieve lines 1-4 from a file at `/path/to/file` you can use a file URI
-  with a region definition such as: `file:///path/to/file:1:0:4:0`.
-
-  When unsure how to fetch a specific region, fall back to fetching the full file
-  contents for a file by omitting the region definition, for example
-  `file:///path/to/file`