Fix imports, annotations, and function naming issues

Co-authored-by: kevinbackhouse <4358136+kevinbackhouse@users.noreply.github.com>

Author: Copilot

src/seclab_taskflows/mcp_servers/alert_results_models.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: 2025 GitHub
 # SPDX-License-Identifier: MIT
 
+from __future__ import annotations
+
 from typing import Optional
 
 from sqlalchemy import Column, ForeignKey, Integer, Text

src/seclab_taskflows/mcp_servers/codeql_python/codeql_sqlite_models.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: 2025 GitHub
 # SPDX-License-Identifier: MIT
 
+from __future__ import annotations
+
 from typing import Optional
 
 from sqlalchemy import Text

src/seclab_taskflows/mcp_servers/codeql_python/mcp_server.py

@@ -1,6 +1,7 @@
 # SPDX-FileCopyrightText: 2025 GitHub
 # SPDX-License-Identifier: MIT
 
+from __future__ import annotations
 
 import csv
 import importlib.resources
@@ -18,8 +19,8 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
-from ..utils import process_repo
-from .codeql_sqlite_models import Base, Source
+from seclab_taskflows.mcp_servers.utils import process_repo
+from seclab_taskflows.mcp_servers.codeql_python.codeql_sqlite_models import Base, Source
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -53,27 +54,29 @@ def source_to_dict(result):
     }
 
 def _resolve_query_path(language: str, query: str) -> Path:
-    global TEMPLATED_QUERY_PATHS
     if language not in TEMPLATED_QUERY_PATHS:
-        raise RuntimeError(f"Error: Language `{language}` not supported!")
+        msg = f"Error: Language `{language}` not supported!"
+        raise RuntimeError(msg)
     query_path = TEMPLATED_QUERY_PATHS[language].get(query)
     if not query_path:
-        raise RuntimeError(f"Error: query `{query}` not supported for `{language}`!")
+        msg = f"Error: query `{query}` not supported for `{language}`!"
+        raise RuntimeError(msg)
     return Path(query_path)
 
 
 def _resolve_db_path(relative_db_path: str | Path):
-    global CODEQL_DBS_BASE_PATH
     # path joins will return "/B" if "/A" / "////B" etc. as well
     # 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).resolve()
     if not absolute_path.is_relative_to(CODEQL_DBS_BASE_PATH.resolve()):
-        raise RuntimeError(f"Error: Database path {absolute_path} is outside the base path {CODEQL_DBS_BASE_PATH}")
+        msg = f"Error: Database path {absolute_path} is outside the base path {CODEQL_DBS_BASE_PATH}"
+        raise RuntimeError(msg)
     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}!")
+        msg = f"Error: Database not found at {absolute_path}!"
+        raise RuntimeError(msg)
     return str(absolute_path)
 
 # This sqlite database is specifically made for CodeQL for Python MCP.
@@ -88,7 +91,7 @@ def __init__(self, memcache_state_dir: str):
         Base.metadata.create_all(self.engine, tables=[Source.__table__])
 
 
-    def store_new_source(self, repo, source_location, line, source_type, notes, update = False):
+    def store_new_source(self, repo, source_location, line, source_type, notes, *, update = False):
         with Session(self.engine) as session:
             existing = session.query(Source).filter_by(repo = repo, source_location = source_location, line = line).first()
             if existing:

src/seclab_taskflows/mcp_servers/gh_code_scanning.py

@@ -1,6 +1,8 @@
 # SPDX-FileCopyrightText: 2025 GitHub
 # SPDX-License-Identifier: MIT
 
+from __future__ import annotations
+
 import json
 import logging
 import os
@@ -17,7 +19,7 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import Session
 
-from .alert_results_models import AlertFlowGraph, AlertResults, Base
+from seclab_taskflows.mcp_servers.alert_results_models import AlertFlowGraph, AlertResults, Base
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -30,6 +32,9 @@
 
 GH_TOKEN = os.getenv('GH_TOKEN', default='')
 
+# Minimum number of parts in an HTML URL to extract owner/repo
+MIN_URL_PARTS = 5
+
 CODEQL_DBS_BASE_PATH = mcp_data_dir('seclab-taskflows', 'codeql', 'CODEQL_DBS_BASE_PATH')
 ALERT_RESULTS_DIR = mcp_data_dir('seclab-taskflows', 'gh_code_scanning', 'ALERT_RESULTS_DIR')
 
@@ -54,7 +59,7 @@ def _get_repo_from_html_url(html_url: str) -> str:
         if not html_url:
             return ''
         parts = html_url.split('/')
-        if len(parts) < 5:
+        if len(parts) < MIN_URL_PARTS:
             return ''
         return f"{parts[3]}/{parts[4]}".lower()
 

src/seclab_taskflows/mcp_servers/ghsa.py

@@ -1,3 +1,8 @@
+# SPDX-FileCopyrightText: 2025 GitHub
+# SPDX-License-Identifier: MIT
+
+from __future__ import annotations
+
 import json
 import logging
 import re
@@ -7,7 +12,9 @@
 from pydantic import Field
 from seclab_taskflow_agent.path_utils import log_file_name
 
-from .gh_code_scanning import call_api
+from seclab_taskflows.mcp_servers.gh_code_scanning import call_api
+
+logger = logging.getLogger(__name__)
 
 logging.basicConfig(
     level=logging.DEBUG,
@@ -22,7 +29,7 @@
 # The advisories contain a lot of information, so we need to filter
 # some of it out to avoid exceeding the maximum prompt size.
 def parse_advisory(advisory: dict) -> dict:
-    logging.debug(f"advisory: {advisory}")
+    logger.debug("advisory: %s", advisory)
     return {
         "ghsa_id": advisory.get("ghsa_id", ""),
         "cve_id": advisory.get("cve_id", ""),
@@ -31,7 +38,7 @@ def parse_advisory(advisory: dict) -> dict:
         "state": advisory.get("state", ""),
     }
 
-async def fetch_GHSA_list_from_gh(owner: str, repo: str) -> str | list:
+async def fetch_ghsa_list_from_gh(owner: str, repo: str) -> str | list:
     """Fetch all security advisories for a specific repository."""
     url = f"https://api.github.com/repos/{owner}/{repo}/security-advisories"
     params = {'per_page': 100}
@@ -60,16 +67,16 @@ async def fetch_GHSA_list_from_gh(owner: str, repo: str) -> str | list:
     return "No advisories found."
 
 @mcp.tool()
-async def fetch_GHSA_list(owner: str = Field(description="The owner of the repo"),
+async def fetch_ghsa_list(owner: str = Field(description="The owner of the repo"),
                           repo: str = Field(description="The repository name")) -> str:
     """Fetch all GitHub Security Advisories (GHSAs) for a specific repository."""
-    results = await fetch_GHSA_list_from_gh(owner, repo)
+    results = await fetch_ghsa_list_from_gh(owner, repo)
     if isinstance(results, str):
         return results
     return json.dumps(results, indent=2)
 
 
-async def fetch_GHSA_details_from_gh(owner: str, repo: str, ghsa_id: str) -> str | dict:
+async def fetch_ghsa_details_from_gh(owner: str, repo: str, ghsa_id: str) -> str | dict:
     """Fetch the details of a repository security advisory."""
     url = f"https://api.github.com/repos/{owner}/{repo}/security-advisories/{ghsa_id}"
     resp = await call_api(url, {})
@@ -80,11 +87,11 @@ async def fetch_GHSA_details_from_gh(owner: str, repo: str, ghsa_id: str) -> str
     return "Not found."
 
 @mcp.tool()
-async def fetch_GHSA_details(owner: str = Field(description="The owner of the repo"),
+async def fetch_ghsa_details(owner: str = Field(description="The owner of the repo"),
                              repo: str = Field(description="The repository name"),
                              ghsa_id: str = Field(description="The ghsa_id of the advisory")) -> str:
     """Fetch a GitHub Security Advisory for a specific repository and GHSA ID."""
-    results = await fetch_GHSA_details_from_gh(owner, repo, ghsa_id)
+    results = await fetch_ghsa_details_from_gh(owner, repo, ghsa_id)
     if isinstance(results, str):
         return results
     return json.dumps(results, indent=2)