1# --------------------------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for license information.
4# --------------------------------------------------------------------------------------------
5
6import threading
7
8from azure.cli.core.commands.parameters import get_resources_in_subscription
9from azure.cli.core.decorators import Completer
10from ._client_factory import cf_network, cf_graph
11
12
13COMPLETION_TIME_OUT = 10
14
15
16# pylint: disable=inconsistent-return-statements
17@Completer
18def subnet_completion_list(cmd, prefix, namespace, **kwargs):  # pylint: disable=unused-argument
19
20    def worker(cmd, prefix, namespace):  # pylint: disable=unused-argument
21        from msrestazure.tools import parse_resource_id
22        client = cf_network(cmd.cli_ctx)
23        subnets = []
24        vnets = get_resources_in_subscription(cmd.cli_ctx, 'Microsoft.Network/virtualNetworks')
25        if namespace.vnet_name and vnets:
26            vnets = [r for r in vnets if r.name.lower() == namespace.vnet_name.lower()]
27            for vnet in vnets:
28                vnet_resource_group = parse_resource_id(vnet.id)['resource_group']
29                for subnet in client.subnets.list(
30                        resource_group_name=vnet_resource_group,
31                        virtual_network_name=namespace.vnet_name):
32                    subnets.append(subnet.name)
33        return subnets
34
35    return HDInsightCompleter(worker=worker).complete(cmd, prefix, namespace)
36
37
38@Completer
39def cluster_admin_account_completion_list(cmd, prefix, namespace, **kwargs):  # pylint: disable=unused-argument
40
41    def worker(cmd, prefix, namespace):  # pylint: disable=unused-argument
42        client = cf_graph(cmd.cli_ctx)
43        filter_str = "startswith(userPrincipalName, '{0}') or startswith(mail, '{0}')" \
44                     .format(prefix) if prefix else None
45        users = client.users.list(filter=filter_str).advance_page()
46        return [user.mail if ("#EXT#" in user.user_principal_name) else user.user_principal_name for user in users]
47
48    return HDInsightCompleter(worker=worker).complete(cmd, prefix, namespace)
49
50
51@Completer
52def cluster_user_group_completion_list(cmd, prefix, namespace, **kwargs):  # pylint: disable=unused-argument
53
54    def worker(cmd, prefix, namespace):  # pylint: disable=unused-argument
55        client = cf_graph(cmd.cli_ctx)
56        filter_str = "startswith(displayName, '{}')".format(prefix) if prefix else None
57        groups = client.groups.list(filter=filter_str).advance_page()
58        return [group.display_name for group in groups]
59
60    return HDInsightCompleter(worker=worker).complete(cmd, prefix, namespace)
61
62
63def get_resource_name_completion_list_under_subscription(resource_type):
64
65    @Completer
66    def completer(cmd, prefix, namespace, **kwargs):  # pylint: disable=unused-argument
67        return [r.name for r in get_resources_in_subscription(cmd.cli_ctx, resource_type)]
68
69    return completer
70
71
72class HDInsightCompleter():
73
74    def __init__(self, worker=None, timeout=COMPLETION_TIME_OUT):
75        self.worker = worker
76        self.timeout = timeout
77        self.worker_result = []
78
79    def complete(self, cmd, prefix, namespace):  # pylint: disable=unused-argument
80        thread = threading.Thread(target=self.complete_worker, args=[self.worker, cmd, prefix, namespace])
81        thread.daemon = True
82        thread.start()
83        thread.join(timeout=self.timeout)
84        return self.worker_result
85
86    def complete_worker(self, worker, cmd, prefix, namespace):
87        self.worker_result = worker(cmd, prefix, namespace)
88