initial public release

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

@@ -0,0 +1,206 @@
+#!/usr/bin/env python
+#
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+from typing import Iterator, List, Optional
+
+from azure.devops.connection import Connection
+from azure.devops.credentials import BasicAuthentication
+from azure.devops.exceptions import AzureDevOpsServiceError
+from azure.devops.v6_0.work_item_tracking.models import (
+    CommentCreate,
+    JsonPatchOperation,
+    Wiql,
+    WorkItem,
+)
+from azure.devops.v6_0.work_item_tracking.work_item_tracking_client import (
+    WorkItemTrackingClient,
+)
+from memoization import cached
+from onefuzztypes.models import ADOTemplate, Report
+
+from .common import Render
+
+
+@cached(ttl=60)
+def get_ado_client(base_url: str, token: str) -> WorkItemTrackingClient:
+    connection = Connection(base_url=base_url, creds=BasicAuthentication("PAT", token))
+    client = connection.clients_v6_0.get_work_item_tracking_client()
+    return client
+
+
+@cached(ttl=60)
+def get_valid_fields(
+    client: WorkItemTrackingClient, project: Optional[str] = None
+) -> List[str]:
+    valid_fields = [
+        x.reference_name.lower()
+        for x in client.get_fields(project=project, expand="ExtensionFields")
+    ]
+    return valid_fields
+
+
+class ADO:
+    def __init__(
+        self, container: str, filename: str, config: ADOTemplate, report: Report
+    ):
+        self.config = config
+        self.renderer = Render(container, filename, report)
+        self.client = get_ado_client(self.config.base_url, self.config.auth_token)
+        self.project = self.render(self.config.project)
+
+    def render(self, template: str) -> str:
+        return self.renderer.render(template)
+
+    def existing_work_items(self) -> Iterator[WorkItem]:
+        filters = {}
+        for key in self.config.unique_fields:
+            if key == "System.TeamProject":
+                value = self.render(self.config.project)
+            else:
+                value = self.render(self.config.ado_fields[key])
+            filters[key.lower()] = value
+
+        valid_fields = get_valid_fields(
+            self.client, project=filters.get("system.teamproject")
+        )
+
+        post_query_filter = {}
+
+        # WIQL (Work Item Query Language) is an SQL like query language that
+        # doesn't support query params, safe quoting, or any other SQL-injection
+        # protection mechanisms.
+        #
+        # As such, build the WIQL with a those fields we can pre-determine are
+        # "safe" and otherwise use post-query filtering.
+        parts = []
+        for k, v in filters.items():
+            # Only add pre-system approved fields to the query
+            if k not in valid_fields:
+                post_query_filter[k] = v
+                continue
+
+            # WIQL supports wrapping values in ' or " and escaping ' by doubling it
+            #
+            # For this System.Title: hi'there
+            # use this query fragment: [System.Title] = 'hi''there'
+            #
+            # For this System.Title: hi"there
+            # use this query fragment: [System.Title] = 'hi"there'
+            #
+            # For this System.Title: hi'"there
+            # use this query fragment: [System.Title] = 'hi''"there'
+            SINGLE = "'"
+            parts.append("[%s] = '%s'" % (k, v.replace(SINGLE, SINGLE + SINGLE)))
+
+        query = "select [System.Id] from WorkItems"
+        if parts:
+            query += " where " + " AND ".join(parts)
+
+        wiql = Wiql(query=query)
+        for entry in self.client.query_by_wiql(wiql).work_items:
+            item = self.client.get_work_item(entry.id, expand="Fields")
+            lowered_fields = {x.lower(): str(y) for (x, y) in item.fields.items()}
+            if post_query_filter and not all(
+                [
+                    k.lower() in lowered_fields and lowered_fields[k.lower()] == v
+                    for (k, v) in post_query_filter.items()
+                ]
+            ):
+                continue
+            yield item
+
+    def update_existing(self, item: WorkItem) -> None:
+        if self.config.on_duplicate.comment:
+            comment = self.render(self.config.on_duplicate.comment)
+            self.client.add_comment(
+                CommentCreate(comment),
+                self.project,
+                item.id,
+            )
+
+        document = []
+        for field in self.config.on_duplicate.increment:
+            value = int(item.fields[field]) if field in item.fields else 0
+            value += 1
+            document.append(
+                JsonPatchOperation(
+                    op="Replace", path="/fields/%s" % field, value=str(value)
+                )
+            )
+
+        for field in self.config.on_duplicate.ado_fields:
+            field_value = self.render(self.config.on_duplicate.ado_fields[field])
+            document.append(
+                JsonPatchOperation(
+                    op="Replace", path="/fields/%s" % field, value=field_value
+                )
+            )
+
+        if item.fields["System.State"] in self.config.on_duplicate.set_state:
+            document.append(
+                JsonPatchOperation(
+                    op="Replace",
+                    path="/fields/System.State",
+                    value=self.config.on_duplicate.set_state[
+                        item.fields["System.State"]
+                    ],
+                )
+            )
+
+        if document:
+            self.client.update_work_item(document, item.id, project=self.project)
+
+    def create_new(self) -> None:
+        task_type = self.render(self.config.type)
+
+        document = []
+        if "System.Tags" not in self.config.ado_fields:
+            document.append(
+                JsonPatchOperation(
+                    op="Add", path="/fields/System.Tags", value="Onefuzz"
+                )
+            )
+
+        for field in self.config.ado_fields:
+            value = self.render(self.config.ado_fields[field])
+            if field == "System.Tags":
+                value += ";Onefuzz"
+            document.append(
+                JsonPatchOperation(op="Add", path="/fields/%s" % field, value=value)
+            )
+
+        entry = self.client.create_work_item(
+            document=document, project=self.project, type=task_type
+        )
+
+        if self.config.comment:
+            comment = self.render(self.config.comment)
+            self.client.add_comment(
+                CommentCreate(comment),
+                self.project,
+                entry.id,
+            )
+
+    def process(self) -> None:
+        seen = False
+        for work_item in self.existing_work_items():
+            self.update_existing(work_item)
+            seen = True
+
+        if not seen:
+            self.create_new()
+
+
+def notify_ado(
+    config: ADOTemplate, container: str, filename: str, report: Report
+) -> None:
+    try:
+        ado = ADO(container, filename, config, report)
+        ado.process()
+    except AzureDevOpsServiceError as err:
+        logging.error("ADO report failed: %s", err)
+    except ValueError as err:
+        logging.error("ADO report value error: %s", err)