MyPy

controller_function/controller/logutils.py

@@ -1,5 +1,6 @@
 """Utils to send logs to Azure Application Insights."""
 import logging
+from typing import Optional
 
 from opencensus.ext.azure.log_exporter import AzureLogHandler
 
@@ -9,7 +10,7 @@
 class CustomDimensionsFilter(logging.Filter):
     """Add application-wide properties to AzureLogHandler records."""
 
-    def __init__(self, custom_dimensions: dict = None) -> None:
+    def __init__(self, custom_dimensions: Optional[dict] = None) -> None:
         """Add custom dimensions, if provided, to the log record."""
         super().__init__()
         self.custom_dimensions = custom_dimensions or {}

controller_function/controller/settings.py

@@ -39,4 +39,4 @@ def correct_start_and_end(cls, v):  # pylint: disable=no-self-argument
 @lru_cache()
 def get_settings() -> Settings:
     """Get the global app settings object."""
-    return Settings()
+    return Settings()  # type: ignore

controller_function/controller/subscription.py

@@ -49,7 +49,11 @@ def enable_subscription(subscription_id: UUID) -> None:
     try:
         SUBSCRIPTION_CLIENT.subscription.enable(str(subscription_id))
     except azure_exceptions.HttpResponseError as e:
-        is_enabled = "not in suspended state" in e.error.message
+        is_enabled = (
+            ("not in suspended state" in e.error.message)
+            if (e.error and e.error.message)
+            else False
+        )
 
         # It's fine if we can't enable it because it is already active.
         if is_enabled:
@@ -70,7 +74,11 @@ def disable_subscription(subscription_id: UUID) -> None:
     try:
         SUBSCRIPTION_CLIENT.subscription.cancel(str(subscription_id))
     except azure_exceptions.HttpResponseError as e:
-        is_disabled = "Subscription is not in active state" in e.error.message
+        is_disabled = (
+            ("Subscription is not in active state" in e.error.message)
+            if (e.error and e.error.message)
+            else False
+        )
 
         # It's fine if we can't disable it because it is already inactive.
         if is_disabled:

controller_function/pyproject.toml

@@ -29,6 +29,8 @@ pytest = "^6.2.5"
 pylint = "^2.17.4"
 pylint-absolute-imports = "^1.0.1"
 isort = "^5.12.0"
+mypy = "^1.8.0"
+types-requests = "^2.31.0.20240125"
 
 [build-system]
 requires = ["poetry-core"]

controller_function/run_tests.sh

@@ -40,6 +40,11 @@ status=$((status+$?))
 python -m coverage report --show-missing
 status=$((status+$?))
 
+# Check types with MyPy
+echo "Running type checking..."
+python -m mypy --config-file tests/mypy.ini controller/ tests/
+status=$((status+$?))
+
 # [optional] Check Markdown coding style with Ruby's markdown lint
 # https://github.com/markdownlint/markdownlint
 if ! command -v mdl &> /dev/null

controller_function/tests/mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = True

status_function/pyproject.toml

@@ -31,6 +31,8 @@ pulumi-azure = "^4.18.0"
 [tool.poetry.group.dev.dependencies]
 pylint-absolute-imports = "^1.0.1"
 isort = "^5.12.0"
+mypy = "^1.8.0"
+types-requests = "^2.31.0.20240125"
 
 [build-system]
 requires = ["poetry-core"]

status_function/run_tests.sh

@@ -38,6 +38,11 @@ status=$((status+$?))
 python -m coverage report --show-missing
 status=$((status+$?))
 
+# Check types with MyPy
+echo "Running type checking..."
+python -m mypy --config-file tests/mypy.ini status/ tests/
+status=$((status+$?))
+
 # [optional] Check Markdown coding style with Ruby's markdown lint
 # https://github.com/markdownlint/markdownlint
 if ! command -v mdl &> /dev/null

status_function/status/__init__.py

@@ -10,6 +10,7 @@
 from datetime import datetime
 from functools import lru_cache
 from typing import Any, Optional
+from uuid import UUID
 
 import azure.functions as func
 import requests
@@ -24,6 +25,7 @@
 from status import models, settings
 from status.auth import BearerAuth
 from status.logutils import set_log_handler
+from status.models import RoleAssignment
 from status.wrapper import CredentialWrapper
 
 logging.basicConfig(
@@ -241,7 +243,7 @@ def get_role_assignment_models(
 def get_subscription_role_assignment_models(
     subscription: Any,
     graph_client: GraphRbacManagementClient,
-) -> list:
+) -> list[RoleAssignment]:
     """Get the role assignment models for each a subscription.
 
     Args:
@@ -259,7 +261,7 @@ def get_subscription_role_assignment_models(
         for assignment in assignments_list:
             role_assignments_models += get_role_assignment_models(
                 assignment,
-                role_def_dict.get(assignment.properties.role_definition_id),
+                role_def_dict.get(assignment.properties.role_definition_id, "Unknown"),
                 graph_client,
             )
     except CloudError as e:
@@ -271,7 +273,7 @@ def get_subscription_role_assignment_models(
     return role_assignments_models
 
 
-def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
+def get_all_status(tenant_id: UUID) -> list[models.SubscriptionStatus]:
     """Get status and role assignments for all subscriptions.
 
     Args:
@@ -284,7 +286,7 @@ def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
     started_at = datetime.now()
 
     graph_client = GraphRbacManagementClient(
-        credentials=GRAPH_CREDENTIALS, tenant_id=tenant_id
+        credentials=GRAPH_CREDENTIALS, tenant_id=str(tenant_id)
     )
 
     client = SubscriptionClient(credential=CREDENTIALS)
@@ -303,7 +305,7 @@ def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
                 subscription_id=subscription.subscription_id,
                 display_name=subscription.display_name,
                 state=subscription.state,
-                role_assignments=role_assignments_models,
+                role_assignments=tuple(role_assignments_models),
             )
         )
 

status_function/status/auth.py

@@ -6,10 +6,15 @@
         value is 5.
 """
 from datetime import datetime, timedelta
+from typing import Any
 
 import jwt
 import requests
 from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import dsa
+from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePrivateKey
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
 from status.settings import get_settings
 
@@ -32,17 +37,24 @@ def __init__(self) -> None:
         # Generate keys with ssh-keygen -t rsa
         private_key_txt = settings.PRIVATE_KEY
 
-        self.private_key = serialization.load_ssh_private_key(
+        private_key = serialization.load_ssh_private_key(
             private_key_txt.encode(), password=b""
         )
+        assert not isinstance(
+            private_key, dsa.DSAPrivateKey
+        ), "DSA keys are not supported."
+
+        # jwt.encode(), used to create the access token, doesn't support DSA keys
+        allowed_type = EllipticCurvePrivateKey | RSAPrivateKey | Ed25519PrivateKey
+        self.private_key: allowed_type = private_key  # type: ignore
 
     def create_access_token(self) -> str:
         """Create an access token for the user to access the API.
 
         Returns:
             The access token.
         """
-        token_claims = {"sub": "status-app"}
+        token_claims: dict[str, Any] = {"sub": "status-app"}
         access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
 
         expire = datetime.utcnow() + access_token_expires

status_function/status/logutils.py

@@ -1,5 +1,6 @@
 """Utils to send logs to Azure Application Insights."""
 import logging
+from typing import Optional
 
 from opencensus.ext.azure.log_exporter import AzureLogHandler
 
@@ -9,7 +10,7 @@
 class CustomDimensionsFilter(logging.Filter):
     """Add application-wide properties to AzureLogHandler records."""
 
-    def __init__(self, custom_dimensions: dict = None) -> None:
+    def __init__(self, custom_dimensions: Optional[dict] = None) -> None:
         """Add custom dimensions to the log record.
 
         Args: