initial public release

src/api-service/__app__/onefuzzlib/notifications/main.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+from typing import Dict, List, Optional, Sequence, Tuple, Union
+from uuid import UUID
+
+from memoization import cached
+from onefuzztypes import models
+from onefuzztypes.enums import ErrorCode, TaskState
+from onefuzztypes.models import ADOTemplate, Error, NotificationTemplate, TeamsTemplate
+from onefuzztypes.primitives import Container, Event
+
+from ..azure.containers import get_container_metadata, get_file_sas_url
+from ..azure.creds import get_fuzz_storage
+from ..azure.queue import send_message
+from ..dashboard import add_event
+from ..orm import ORMMixin
+from ..reports import get_report
+from ..tasks.config import get_input_container_queues
+from ..tasks.main import Task
+from .ado import notify_ado
+from .teams import notify_teams
+
+
+class Notification(models.Notification, ORMMixin):
+    @classmethod
+    def get_by_id(cls, notification_id: UUID) -> Union[Error, "Notification"]:
+        notifications = cls.search(query={"notification_id": [notification_id]})
+        if not notifications:
+            return Error(
+                code=ErrorCode.INVALID_REQUEST, errors=["unable to find Notification"]
+            )
+
+        if len(notifications) != 1:
+            return Error(
+                code=ErrorCode.INVALID_REQUEST,
+                errors=["error identifying Notification"],
+            )
+        notification = notifications[0]
+        return notification
+
+    @classmethod
+    def get_existing(
+        cls, container: Container, config: NotificationTemplate
+    ) -> Optional["Notification"]:
+        notifications = Notification.search(query={"container": [container]})
+        for notification in notifications:
+            if notification.config == config:
+                return notification
+        return None
+
+    @classmethod
+    def key_fields(cls) -> Tuple[str, str]:
+        return ("notification_id", "container")
+
+
+@cached(ttl=10)
+def get_notifications(container: Container) -> List[Notification]:
+    return Notification.search(query={"container": [container]})
+
+
+@cached(ttl=10)
+def get_queue_tasks() -> Sequence[Tuple[Task, Sequence[str]]]:
+    results = []
+    for task in Task.search_states(states=TaskState.available()):
+        containers = get_input_container_queues(task.config)
+        if containers:
+            results.append((task, containers))
+    return results
+
+
+@cached(ttl=60)
+def container_metadata(container: Container) -> Optional[Dict[str, str]]:
+    return get_container_metadata(container)
+
+
+def new_files(container: Container, filename: str) -> None:
+    results: Dict[str, Event] = {"container": container, "file": filename}
+
+    metadata = container_metadata(container)
+    if metadata:
+        results["metadata"] = metadata
+
+    notifications = get_notifications(container)
+    if notifications:
+        report = get_report(container, filename)
+        if report:
+            results["executable"] = report.executable
+            results["crash_type"] = report.crash_type
+            results["crash_site"] = report.crash_site
+            results["job_id"] = report.job_id
+            results["task_id"] = report.task_id
+
+        logging.info("notifications for %s %s %s", container, filename, notifications)
+        done = []
+        for notification in notifications:
+            # ignore duplicate configurations
+            if notification.config in done:
+                continue
+            done.append(notification.config)
+
+            if isinstance(notification.config, TeamsTemplate):
+                notify_teams(notification.config, container, filename, report)
+
+            if not report:
+                continue
+
+            if isinstance(notification.config, ADOTemplate):
+                notify_ado(notification.config, container, filename, report)
+
+    for (task, containers) in get_queue_tasks():
+        if container in containers:
+            logging.info("queuing input %s %s %s", container, filename, task.task_id)
+            url = get_file_sas_url(container, filename, read=True, delete=True)
+            send_message(
+                task.task_id, bytes(url, "utf-8"), account_id=get_fuzz_storage()
+            )
+
+    add_event("new_file", results)