datadogpy/datadog/dogshell/common.py at master · ryanrhee/datadogpy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
# Unless explicitly stated otherwise all files in this repository are licensed under the BSD-3-Clause License.
# This product includes software developed at Datadog (https://www.datadoghq.com/).
# Copyright 2015-Present Datadog, Inc
# stdlib
from __future__ import print_function
import os
import sys

# datadog
from datadog.util.compat import is_p3k, configparser, IterableUserDict, get_input


def print_err(msg):
    if is_p3k():
        print(msg + "\n", file=sys.stderr)
    else:
        sys.stderr.write(msg + "\n")
    sys.stderr.flush()


def report_errors(res):
    if "errors" in res:
        errors = res["errors"]
        if isinstance(errors, list):
            for error in errors:
                print_err("ERROR: {}".format(error))
        else:
            print_err("ERROR: {}".format(errors))
        sys.exit(1)
    return False


def report_warnings(res):
    if "warnings" in res:
        warnings = res["warnings"]
        if isinstance(warnings, list):
            for warning in warnings:
                print_err("WARNING: {}".format(warning))
        else:
            print_err("WARNING: {}".format(warnings))
        return True
    return False


class DogshellConfig(IterableUserDict):
    def load(self, config_file, api_key, app_key, api_host):
        config = configparser.ConfigParser()

        if api_host is not None:
            if api_host in ("us" or "datadoghq.com"):
                self["api_host"] = "https://datadoghq.com"
            elif api_host in ("datadoghq.eu", "eu"):
                self["api_host"] = "https://datadoghq.eu"
            elif api_host in ("us3.datadoghq.com", "us3"):
                self["api_host"] = "https://us3.datadoghq.com"

        if api_key is not None and app_key is not None:
            self["api_key"] = api_key
            self["app_key"] = app_key
        else:
            if os.access(config_file, os.F_OK):
                config.read(config_file)
                if not config.has_section("Connection"):
                    report_errors({"errors": ["%s has no [Connection] section" % config_file]})
            else:
                try:
                    response = None
                    while response is None or response.strip().lower() not in ["", "y", "n"]:
                        response = get_input("%s does not exist. Would you like to" " create it? [Y/n] " % config_file)
                        if response.strip().lower() in ["", "y"]:
                            # Read the api and app keys from stdin
                            api_key = get_input(
                                "What is your api key? (Get it here: "
                                "https://app.datadoghq.com/account/settings#api) "
                            )
                            app_key = get_input(
                                "What is your application key? (Generate one here: "
                                "https://app.datadoghq.com/account/settings#api) "
                            )

                            # Write the config file
                            config.add_section("Connection")
                            config.set("Connection", "apikey", api_key)
                            config.set("Connection", "appkey", app_key)

                            f = open(config_file, "w")
                            config.write(f)
                            f.close()
                            print("Wrote %s" % config_file)
                        elif response.strip().lower() == "n":
                            # Abort
                            print_err("Exiting\n")
                            sys.exit(1)
                except KeyboardInterrupt:
                    # Abort
                    print_err("\nExiting")
                    sys.exit(1)

            self["api_key"] = config.get("Connection", "apikey")
            self["app_key"] = config.get("Connection", "appkey")
            if config.has_section("Proxy"):
                self["proxies"] = dict(config.items("Proxy"))
            if config.has_option("Connection", "host_name"):
                self["host_name"] = config.get("Connection", "host_name")
            if config.has_option("Connection", "api_host"):
                self["api_host"] = config.get("Connection", "api_host")
        assert self["api_key"] is not None and self["app_key"] is not None