1# -*- coding: utf-8 -*-
2
3# This Source Code Form is subject to the terms of the Mozilla Public
4# License, v. 2.0. If a copy of the MPL was not distributed with this
5# file, You can obtain one at http://mozilla.org/MPL/2.0/.
6
7from __future__ import absolute_import, print_function, unicode_literals
8
9import logging
10import os
11import shutil
12import subprocess
13import tempfile
14from contextlib import contextmanager
15
16from pathlib import Path
17
18from appdirs import user_config_dir
19import taskcluster
20
21from mach.base import FailedCommandError
22from taskgraph import GECKO
23
24logger = logging.getLogger(__name__)
25
26
27TASK_TYPES = {
28    "signing": ["linux-signing", "linux-signing-partial"],
29    "beetmover": ["beetmover-candidates"],
30    "bouncer": ["bouncer-submit"],
31    "balrog": ["balrog-submit"],
32    "tree": ["tree"],
33}
34
35
36def get_secret(secret):
37    # use proxy if configured, otherwise local credentials from env vars
38    if "TASKCLUSTER_PROXY_URL" in os.environ:
39        secrets_options = {"rootUrl": os.environ["TASKCLUSTER_PROXY_URL"]}
40    else:
41        secrets_options = taskcluster.optionsFromEnvironment()
42    secrets = taskcluster.Secrets(secrets_options)
43    return secrets.get(secret)["secret"]
44
45
46@contextmanager
47def configure_ssh(ssh_key_secret):
48    if ssh_key_secret is None:
49        yield
50
51    # If we get here, we are runnig in automation.
52    # We use a user hgrc, so that we are also get the system-wide hgrc
53    # settings.
54    hgrc = Path(user_config_dir("hg")).joinpath("hgrc")
55    if hgrc.exists():
56        raise FailedCommandError(
57            "Not overwriting `{}`; cannot configure ssh.".format(hgrc)
58        )
59
60    try:
61        ssh_key_dir = Path(tempfile.mkdtemp())
62
63        ssh_key = get_secret(ssh_key_secret)
64        ssh_key_file = ssh_key_dir.joinpath("id_rsa")
65        ssh_key_file.write_text(ssh_key["ssh_privkey"])
66        ssh_key_file.chmod(0o600)
67
68        hgrc_content = (
69            "[ui]\n"
70            "username = trybld\n"
71            "ssh = ssh -i {path} -l {user}\n".format(
72                path=ssh_key_file, user=ssh_key["user"]
73            )
74        )
75        hgrc.write_text(hgrc_content)
76
77        yield
78    finally:
79        shutil.rmtree(str(ssh_key_dir))
80        os.remove(str(hgrc))
81
82
83def push_canary(scriptworkers, addresses, ssh_key_secret):
84    if ssh_key_secret and os.environ.get("MOZ_AUTOMATION", "0") != "1":
85        # We make assumptions about the layout of the docker image
86        # for creating the hgrc that we use for the key.
87        raise FailedCommandError("Cannot use ssh-key-secret outside of automation.")
88
89    # Collect the set of `mach try scriptworker` task sets to run.
90    tasks = []
91    for scriptworker in scriptworkers:
92        worker_tasks = TASK_TYPES.get(scriptworker)
93        if worker_tasks:
94            logger.info("Running tasks for {}: {}".format(scriptworker, worker_tasks))
95            tasks.extend(worker_tasks)
96        else:
97            logger.info("No tasks for {}.".format(scriptworker))
98
99    mach = Path(GECKO).joinpath("mach")
100    base_command = [str(mach), "try", "scriptworker"]
101    for address in addresses:
102        base_command.extend(
103            [
104                "--route",
105                "notify.email.{}.on-failed".format(address),
106                "--route",
107                "notify.email.{}.on-exception".format(address),
108            ]
109        )
110
111    with configure_ssh(ssh_key_secret):
112        env = os.environ.copy()
113        for task in tasks:
114            subprocess.check_call(base_command + [task], env=env)
115