recipe_modules/tree_status/resources/tree_status.py

#!/usr/bin/env vpython3
# Copyright 2020 The Fuchsia Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# [VPYTHON:BEGIN]
# python_version: "3.8"
# wheel: <
#   name: "infra/python/wheels/pyasn1_modules-py3"
#   version: "version:0.2.8"
# >
# wheel: <
#   name: "infra/python/wheels/cachetools-py3"
#   version: "version:4.2.2"
# >
# wheel: <
#   name: "infra/python/wheels/six-py3"
#   version: "version:1.15.0"
# >
# wheel: <
#   name: "infra/python/wheels/pyasn1-py3"
#   version: "version:0.4.8"
# >
# wheel: <
#   name: "infra/python/wheels/certifi-py3"
#   version: "version:2020.4.5.1"
# >
# wheel: <
#   name: "infra/python/wheels/rsa-py3"
#   version: "version:4.7.2"
# >
# wheel: <
#   name: "infra/python/wheels/idna-py3"
#   version: "version:2.8"
# >
# wheel: <
#   name: "infra/python/wheels/chardet-py3"
#   version: "version:4.0.0"
# >
# wheel: <
#   name: "infra/python/wheels/urllib3-py3"
#   version: "version:1.26.4"
# >
# wheel: <
#   name: "infra/python/wheels/requests-py3"
#   version: "version:2.25.1"
# >
# wheel: <
#   name: "infra/python/wheels/google-auth-py3"
#   version: "version:2.6.0"
# >
# [VPYTHON:END]

"""Get and set the status of the Fuchsia tree.

Fetch tree status (prints results to stdout):
  ./tree_status.py fuchsia.stem-status.appspot.com get
  > {"general_state": "closed", "message": "tree is closed :(", ...}

Update tree status:
  ./tree_status.py fuchsia-stem-status.appspot.com set 'tree is closed' \
    --admin_hostname=fuchsia-status.adm.googleplex.com --username someuser
  > OK
"""

import argparse
import contextlib
import http.client
import os
import requests
import time
from typing import Dict
import urllib.parse
import warnings

JSON_PREFIX = ")]}'"

TREE_STATUS_SERVICE = (
    "https://luci-tree-status.appspot.com/prpc/luci.tree_status.v1.TreeStatus"
)


def main():
    parser = argparse.ArgumentParser(description="Get and set tree status.")
    parser.add_argument("hostname", help="Hostname for the status page")
    parser.add_argument(
        "--tree-name",
        default="fuchsia-stem",
        help="Name of the tree to get the status for",
    )
    parser.add_argument(
        "--use-new-status-app",
        action="store_true",
        help="Whether to use the new tree status app",
    )

    subparsers = parser.add_subparsers(help="Subcommands", required=True)

    get_parser = subparsers.add_parser("get", help="Get the tree status")
    get_parser.set_defaults(func=get_status)

    set_parser = subparsers.add_parser("set", help="Set the tree status")
    set_parser.set_defaults(func=set_status)
    set_parser.add_argument("message", help="Message to set as tree status")
    set_parser.add_argument(
        "--admin_hostname",
        default="fuchsia-status-adm.googleplex.com",
        help="Hostname for admin server",
    )
    set_parser.add_argument(
        "--username", help="Username to use for setting the status", required=True
    )
    set_parser.add_argument(
        "--state", help="State to set (if not inferred from message)"
    )
    set_parser.add_argument("--issuetracker_id", help="IssueTracker ID to filed bug")

    args = parser.parse_args()

    return args.func(args)


def get_status(args):
    max_attempts = 5
    sleep_seconds = 1
    for _ in range(max_attempts):
        try:
            result = make_googleplex_request(
                f"{TREE_STATUS_SERVICE}/GetStatus",
                {"name": f"trees/{args.tree_name}/status/latest"},
            )
            if not args.use_new_status_app:
                result = make_request(args.hostname, "/current?format=json")

        except BadHTTPResponse as e:
            if e.status < 500:
                raise
            # Transient server error, retry with exponential backoff.
            time.sleep(sleep_seconds)
            sleep_seconds *= 2
        else:
            # Success! All done.
            print(result)
            break


def set_status(args):
    try:
        # TODO(https://fxbug.dev/332741591): We can stop setting the status in
        # both apps once we no longer need the old app for the stats dashboards.
        result = make_googleplex_request(
            f"{TREE_STATUS_SERVICE}/CreateStatus",
            {
                "parent": f"trees/{args.tree_name}/status",
                "status": {
                    "generalState": args.state.upper(),
                    "message": args.message,
                },
            },
        )
        old_app_result = make_googleplex_request(
            f"https://{args.admin_hostname}/api/v1/status",
            {
                "user": args.username,
                "message": args.message,
                "state": args.state,
                "issue_tracker_id": args.issuetracker_id,
            },
        )
        if not args.use_new_status_app:
            result = old_app_result
        print(result)
    except Exception as e:
        warnings.warn(str(e))


class BadHTTPResponse(Exception):
    def __init__(self, status, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.status = status


def make_request(url, path, method="GET", params=None):
    # Pylint can't tell that `closing()` returns the same thing as the function
    # that it wraps.
    # pylint: disable=no-member
    with contextlib.closing(http.client.HTTPSConnection(url)) as conn:
        encoded_params = urllib.parse.urlencode(params) if params else ""

        conn.request(method, path, encoded_params)
        response = conn.getresponse()
        payload = response.read().decode()
        if response.status >= 300:
            raise BadHTTPResponse(
                int(response.status),
                f"bad response status {int(response.status)}: {response.reason} with payload `{payload}`",
            )
        return payload
    # pylint: enable=no-member


def make_googleplex_request(url: str, json: Dict[str, str]) -> str:
    tokenResp = requests.get(
        f"http://{os.getenv('GCE_METADATA_HOST')}/computeMetadata/v1/instance/service-accounts/default/token",
        headers={"Metadata-Flavor": "Google"},
    )

    resp = requests.post(
        url,
        headers={
            "Accept": "application/json",
            "Authorization": "Bearer " + tokenResp.json()["access_token"],
            "Content-Type": "application/json",
        },
        json=json,
        # Allowing redirects can hide authorization issues.
        # Disallow them to make it clear when there is a 403 Unauthorized.
        allow_redirects=False,
    )
    if resp.status_code >= 500:
        raise Exception(f"bad response status {int(resp.status)}: {resp.text}")

    result = resp.text
    if result.startswith(JSON_PREFIX):
        result = result[len(JSON_PREFIX) :]

    return result


if __name__ == "__main__":
    main()