seantis · somehowchris · Jan 28, 2025 · Jan 28, 2025
diff --git a/src/riskmatrix/__init__.py b/src/riskmatrix/__init__.py
@@ -12,8 +12,6 @@
 from riskmatrix.route_factories import root_factory
 from riskmatrix.security import authenticated_user
 from riskmatrix.security_policy import SessionSecurityPolicy
-from openai import OpenAI
-from anthropic import Anthropic
 from langchain_openai import ChatOpenAI
 from langchain_anthropic import ChatAnthropic
 
@@ -95,7 +93,7 @@ def main(
         if openai_apikey := settings.get('openai_api_key'):
             openai_client = ChatOpenAI(
                 api_key=openai_apikey,
-                model = "gpt-4o-mini",
+                model="gpt-4o-mini",
                 temperature=0.7
             )
             config.add_request_method(
@@ -114,7 +112,7 @@ def main(
                 'llm',
                 reify=True
             )
-            
+
         if langfuse_host := settings.get("langfuse_host"):
             from langfuse.callback import CallbackHandler
             langfuse_handler = CallbackHandler(
@@ -127,7 +125,6 @@ def main(
                 'langfuse',
                 reify=True
             )
-
 
     app = config.make_wsgi_app()
     return Fanstatic(app, versioning=True)
diff --git a/src/riskmatrix/layouts/layout.pt b/src/riskmatrix/layouts/layout.pt
@@ -5,52 +5,45 @@
       i18n:domain="riskmatrix"
       lang="${layout.locale_name()}"
       class="h-100"
-      data_sentry_dsn="${layout.sentry_dsn()}">
+      data_sentry_dsn="${layout.sentry_dsn()}"
+      data-bs-theme="dark">
 
     <head>
-
         <meta charset="utf-8" />
         <meta name="viewport" content="width=device-width, initial-scale=1" />
 
         <tal:block tal:condition="layout.sentry_dsn()">
-            <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/bundle.min.js')}"></script>
-            <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/sentry.js')}"></script>
+            <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/bundle.min.js')}" nonce="${request.csp_nonce}"></script>
+            <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/sentry.js')}" nonce="${request.csp_nonce}"></script>
         </tal:block>
-        <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/plotly.min.js')}"></script>
-        <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/marked.min.js')}"></script>
-
-        <title>RiskMatrix<tal:b tal:condition="exists:title"> — ${title}</tal:b></title>
+        <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/plotly.min.js')}" nonce="${request.csp_nonce}"></script>
+        <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/marked.min.js')}" nonce="${request.csp_nonce}"></script>
+        <script type="text/javascript" src="${layout.static_url('riskmatrix:static/js/plotly_theme.js')}" nonce="${request.csp_nonce}"></script>
 
+        <title>RiskMatrix<tal:b tal:condition="exists:title"> — ${title}</tal:b></title>
     </head>
 
     <body class="d-flex flex-column h-100">
-
+    
         ${panel('navbar')}
 
         <!-- Begin page content -->
         <main class="flex-shrink-0">
-
             ${panel('flash')}
-
             <div class="container">
-
                 <div class="row" tal:omit-tag="not:layout.show_steps()">
-
                     <div class="col-md-auto" tal:condition="layout.show_steps()">
                         ${panel('steps')}
                     </div>
-
                     <div class="col" tal:omit-tag="not:layout.show_steps()">
                         <tal:block metal:define-slot="content" />
                     </div>
                 </div>
-
             </div>
-
             <tal:block metal:define-slot="modals" />
-
         </main>
 
+        <!-- Footer Content -->
         <footer class="footer mt-auto py-3 bg-dark">
             <div class="container">
                 <div class="row">
@@ -108,7 +101,5 @@
                 </div>
             </div>
         </footer>
-
     </body>
-
 </html>
diff --git a/src/riskmatrix/layouts/navbar.pt b/src/riskmatrix/layouts/navbar.pt
@@ -12,6 +12,36 @@
                     </li>
                 </ul>
                 <a class="btn btn-outline-light" href="${request.route_url('logout')}"><i class="fad fa-sign-out"></i> <tal:b i18n:translate>Logout</tal:b></a>
+                &nbsp;
+                <div class="dropdown">
+                    <button class="btn btn-link nav-link py-2 px-0 px-lg-2 dropdown-toggle d-flex align-items-center" id="bd-theme" type="button" aria-expanded="false" data-bs-toggle="dropdown" data-bs-display="static" aria-label="Toggle theme (light)">
+                        <i class="fa fa-sun theme-icon-active" aria-hidden="true"></i>
+                        <span class="d-lg-none ms-2" id="bd-theme-text">Toggle theme</span>
+                    </button>
+                    <ul class="dropdown-menu dropdown-menu-end" aria-labelledby="bd-theme-text">
+                        <li>
+                            <button type="button" class="dropdown-item d-flex align-items-center" data-bs-theme-value="light" aria-pressed="true">
+                                <i class="fa fa-sun" aria-hidden="true"></i>
+                                &nbsp;
+                                Light
+                            </button>
+                        </li>
+                        <li>
+                            <button type="button" class="dropdown-item d-flex align-items-center" data-bs-theme-value="dark" aria-pressed="false">
+                                <i class="fa fa-moon" aria-hidden="true"></i>
+                                &nbsp;
+                                Dark
+                            </button>
+                        </li>
+                        <li>
+                            <button type="button" class="dropdown-item d-flex align-items-center" data-bs-theme-value="auto" aria-pressed="false">
+                                <i class="fa fa-laptop" aria-hidden="true"></i>
+                                &nbsp;
+                                Auto
+                            </button>
+                        </li>
+                    </ul>
+                </div>
             </div>
         </div>
     </nav>

diff --git a/src/riskmatrix/layouts/navbar.py b/src/riskmatrix/layouts/navbar.py
@@ -67,7 +67,9 @@ def navbar(context: object, request: 'IRequest') -> 'RenderData':
                 request,
                 _('Risk Catalog'),
                 request.route_url('risk_catalog'),
-                lambda request, url: request.path_url.startswith(request.route_url('risk_catalog'))
+                lambda request, url: request.path_url.startswith(
+                    request.route_url('risk_catalog')
+                )
             ),
             NavbarEntry(
                 request,

diff --git a/src/riskmatrix/layouts/steps.py b/src/riskmatrix/layouts/steps.py
@@ -4,8 +4,6 @@
 from typing import NamedTuple
 from typing import TYPE_CHECKING
 
-from riskmatrix.models.risk_assessment import RiskAssessment
-from riskmatrix.models.risk_assessment_info import RiskAssessmentInfo, RiskAssessmentState
 if TYPE_CHECKING:
     from pyramid.interfaces import IRequest
 
@@ -20,12 +18,6 @@ class Step(NamedTuple):
 
 
 def steps(context: 'Organization', request: 'IRequest') -> 'RenderData':
-    assessments = request.dbsession.query(RiskAssessmentInfo).filter(
-        RiskAssessmentInfo.organization_id == context.id,
-        RiskAssessmentInfo.state == RiskAssessmentState.OPEN,
-    ).all()
-
-    t = request.dbsession.query(RiskAssessment).filter(RiskAssessment.risk_assessment_info_id.in_([a.id for a in assessments]), RiskAssessment.likelihood == None, RiskAssessment.impact == None).count()
     return {
         'steps': [
             Step(
@@ -49,7 +41,11 @@ def steps(context: 'Organization', request: 'IRequest') -> 'RenderData':
                 '#',
                 disabled=True
             ),
-            Step(_("Finish Assessment"), request.route_url('finish_assessment'), disabled=False),#t > 0 or len(assessments) == 0),
+            Step(
+                _("Finish Assessment"),
+                request.route_url('finish_assessment'),
+                disabled=False
+            ),
         ]
     }
 

diff --git a/src/riskmatrix/mail/__init__.py b/src/riskmatrix/mail/__init__.py
@@ -12,4 +12,4 @@
     'MailError',
     'MailState',
     'PostmarkMailer',
-)
+)
diff --git a/src/riskmatrix/mail/exceptions.py b/src/riskmatrix/mail/exceptions.py
@@ -7,4 +7,4 @@ class MailConnectionError(MailError, ConnectionError):
 
 
 class InactiveRecipient(MailError):
-    pass
+    pass
diff --git a/src/riskmatrix/mail/interfaces.py b/src/riskmatrix/mail/interfaces.py
@@ -88,4 +88,4 @@ def delete_template(template: 'ITemplate') -> list[str]:
         Deletes a mailer template based on a certificate template.
 
         Returns a list of errors. If the list is empty, it was successful.
-        """
+        """
diff --git a/src/riskmatrix/mail/mailer.py b/src/riskmatrix/mail/mailer.py
@@ -547,8 +547,12 @@ def create_or_update_template(
             name = f'[{self.stream}] {organization.name}: {template.name}'
         if len(name) > 100:
             name = name[:97] + '...'
-        html_content: str = markdown_to_html(template.email_content)
-        plain_content = markdown_to_plaintext(template.email_content)
+        html_content: str = markdown_to_html(  # noqa: F821
+            template.email_content
+        )
+        plain_content = markdown_to_plaintext(  # noqa: F821
+            template.email_content
+        )
 
         # replace logos with appropriate placeholders.
         html_content = html_content.replace(
@@ -622,4 +626,4 @@ def delete_template(self, template: 'ITemplate') -> list[str]:
     'Transient': 0,  # Needs to have lower priority than hard bounce
     'Bounced': MailState.failed,  # This is a hard bounce
     'Opened': MailState.read,
-}
+}
diff --git a/src/riskmatrix/mail/types.py b/src/riskmatrix/mail/types.py
@@ -64,4 +64,4 @@ class TemplateMailParams(
     RequiredTemplateMailParams,
     OptionalTemplateMailParams
 ):
-    pass
+    pass
diff --git a/src/riskmatrix/models/asset.py b/src/riskmatrix/models/asset.py
@@ -28,6 +28,7 @@
 
 from sqlalchemy_serializer import SerializerMixin
 
+
 class Asset(Base, SoftDeleteMixin, SerializerMixin):
 
     __tablename__ = 'asset'
@@ -50,7 +51,7 @@ class Asset(Base, SoftDeleteMixin, SerializerMixin):
     assessments: Mapped[list['RiskAssessment']] = relationship(
         back_populates='asset'
     )
-    
+
     organization: Mapped['Organization'] = relationship(
         back_populates='assets'
     )

diff --git a/src/riskmatrix/models/organization.py b/src/riskmatrix/models/organization.py
@@ -33,7 +33,7 @@ class Organization(Base):
     risks: Mapped[list['Risk']] = relationship(
         back_populates='organization',
     )
-    risk_catalogs:  Mapped[list['RiskCatalog']] = relationship(
+    risk_catalogs:  Mapped[list[RiskCatalog]] = relationship(
         back_populates='organization',
     )
     users: Mapped[list['User']] = relationship(

diff --git a/src/riskmatrix/models/password_change_token.py b/src/riskmatrix/models/password_change_token.py
@@ -98,4 +98,4 @@ def expired(self) -> bool:
         expiring_time = self.time_requested + timedelta(hours=48)
         if datetime.utcnow() > expiring_time:
             return True
-        return False
+        return False
diff --git a/src/riskmatrix/models/risk.py b/src/riskmatrix/models/risk.py
@@ -28,6 +28,7 @@
 
 from sqlalchemy_serializer import SerializerMixin
 
+
 class Risk(SoftDeleteMixin, Base, SerializerMixin):
 
     __tablename__ = 'risk'

diff --git a/src/riskmatrix/models/risk_assessment.py b/src/riskmatrix/models/risk_assessment.py
@@ -18,7 +18,6 @@
 from riskmatrix.orm.meta import UUIDStr
 from riskmatrix.orm.meta import UUIDStrPK
 from sqlalchemy.types import JSON
-from dataclasses import dataclass
 from sqlalchemy_serializer import SerializerMixin
 
 from typing import Any, ClassVar
@@ -82,7 +81,7 @@ def __init__(
         self,
         asset:  Asset,
         risk:   Risk,
-        info:  'RiskAssessmentInfo',
+        info:  RiskAssessmentInfo,
         **meta: Any
     ):
         self.id = str(uuid4())

diff --git a/src/riskmatrix/models/risk_assessment_info.py b/src/riskmatrix/models/risk_assessment_info.py
@@ -1,4 +1,3 @@
-from dataclasses import dataclass
 from sedate import utcnow
 from riskmatrix.orm.meta import Base
 import enum
@@ -8,16 +7,13 @@
 from sqlalchemy.orm import Mapped
 from riskmatrix.orm.meta import UUIDStr
 from riskmatrix.orm.meta import UUIDStrPK
-from sqlalchemy import UniqueConstraint
-from sqlalchemy import Column
+from sqlalchemy_serializer import SerializerMixin
 from datetime import datetime
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from riskmatrix.models import RiskAssessment
-
-
 
 
 class RiskAssessmentState(enum.Enum):
@@ -31,8 +27,6 @@ def __str__(self) -> str:
         }
         return names[self]
 
-from sqlalchemy_serializer import SerializerMixin
-
 
 class RiskAssessmentInfo(Base, SerializerMixin):
 
@@ -43,7 +37,7 @@ class RiskAssessmentInfo(Base, SerializerMixin):
         ForeignKey('organization.id', ondelete='CASCADE'),
         index=True,
     )
-    
+
     name: Mapped[str] = mapped_column(nullable=True)
 
     state: Mapped[RiskAssessmentState] = mapped_column(

diff --git a/src/riskmatrix/models/risk_catalog.py b/src/riskmatrix/models/risk_catalog.py
@@ -58,7 +58,9 @@ class RiskCatalog(Base, SoftDeleteMixin, SerializerMixin):
     modified: Mapped[datetime | None] = mapped_column(onupdate=utcnow)
 
     risks: Mapped[list['Risk']] = relationship(back_populates='catalog')
-    organization: Mapped['Organization'] = relationship(back_populates='risk_catalogs')
+    organization: Mapped['Organization'] = relationship(
+        back_populates='risk_catalogs'
+    )
 
     def __init__(
         self,

diff --git a/src/riskmatrix/models/risk_treatment.py b/src/riskmatrix/models/risk_treatment.py
diff --git a/src/riskmatrix/orm/softdelete_base.py b/src/riskmatrix/orm/softdelete_base.py
@@ -1,18 +1,17 @@
 from sqlalchemy_easy_softdelete.mixin import generate_soft_delete_mixin_class
-from sqlalchemy_easy_softdelete.hook import IgnoredTable
-from sqlalchemy.orm import declarative_base
-from sqlalchemy import Column, Integer
-from datetime import datetime
+
 from sqlalchemy.sql import func
 
-# Create a Class that inherits from our class builder
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from datetime import datetime
+
+
 class SoftDeleteMixin(generate_soft_delete_mixin_class(
-    # This table will be ignored by the hook
-    # even if the table has the soft-delete column
     ignored_tables=[]
 )):
-    # type hint for autocomplete IDE support
-    deleted_at: datetime
+    deleted_at: 'datetime'
 
     def soft_delete(self):
-        self.deleted_at = func.now()
+        self.deleted_at = func.now()
diff --git a/src/riskmatrix/prompts/__init__.py b/src/riskmatrix/prompts/__init__.py
@@ -0,0 +1,17 @@
+import json
+from pathlib import Path
+
+
+def load_json_file(filename: str) -> dict:
+    """Load a JSON file from the prompts directory."""
+    current_dir = Path(__file__).parent
+    file_path = current_dir / filename
+
+    with open(file_path, 'r', encoding='utf-8') as f:
+        return json.load(f)
+
+
+# Load all prompt configurations
+system_prompts = load_json_file('system_prompts.json')
+user_prompts = load_json_file('user_prompts.json')
+examples = load_json_file('examples.json')
Original file line number	Diff line number	Diff line change
Expand Up		@@ -7,4 +7,4 @@ class MailConnectionError(MailError, ConnectionError):


		class InactiveRecipient(MailError):
		pass
		pass
-Original file line number
+Diff line change
@@ Expand Up / @@ -88,4 +88,4 @@ def delete_template(template: 'ITemplate') -> list[str]: @@
             Deletes a mailer template based on a certificate template.
             Returns a list of errors. If the list is empty, it was successful.
-            """
+            """
Original file line number	Diff line number	Diff line change
Expand Up		@@ -28,6 +28,7 @@

		from sqlalchemy_serializer import SerializerMixin


		class Risk(SoftDeleteMixin, Base, SerializerMixin):

		__tablename__ = 'risk'
Expand Down