1import os
2from typing import TypeVar
3from typing import Union
4
5from ddtrace.internal.compat import parse
6from ddtrace.internal.utils.formats import get_env
7
8from .http import HTTPConnection
9from .http import HTTPSConnection
10from .uds import UDSHTTPConnection
11
12
13DEFAULT_HOSTNAME = "localhost"
14DEFAULT_TRACE_PORT = 8126
15DEFAULT_UNIX_TRACE_PATH = "/var/run/datadog/apm.socket"
16DEFAULT_UNIX_DSD_PATH = "/var/run/datadog/dsd.socket"
17DEFAULT_STATS_PORT = 8125
18DEFAULT_TRACE_URL = "http://%s:%s" % (DEFAULT_HOSTNAME, DEFAULT_TRACE_PORT)
19DEFAULT_TIMEOUT = 2.0
20
21ConnectionType = Union[HTTPSConnection, HTTPConnection, UDSHTTPConnection]
22
23T = TypeVar("T")
24
25
26def get_trace_hostname(default=DEFAULT_HOSTNAME):
27    # type: (Union[T, str]) -> Union[T, str]
28    return os.environ.get("DD_AGENT_HOST", os.environ.get("DATADOG_TRACE_AGENT_HOSTNAME", default))
29
30
31def get_stats_hostname(default=DEFAULT_HOSTNAME):
32    # type: (Union[T, str]) -> Union[T, str]
33    return os.environ.get("DD_AGENT_HOST", os.environ.get("DD_DOGSTATSD_HOST", default))
34
35
36def get_trace_port(default=DEFAULT_TRACE_PORT):
37    # type: (Union[T, int]) -> Union[T,int]
38    v = os.environ.get("DD_AGENT_PORT", os.environ.get("DD_TRACE_AGENT_PORT"))
39    if v is not None:
40        return int(v)
41    return default
42
43
44def get_stats_port(default=DEFAULT_STATS_PORT):
45    # type: (Union[T, int]) -> Union[T,int]
46    v = get_env("dogstatsd", "port", default=None)
47    if v is not None:
48        return int(v)
49    return default
50
51
52def get_trace_agent_timeout():
53    # type: () -> float
54    return float(get_env("trace", "agent", "timeout", "seconds", default=DEFAULT_TIMEOUT))  # type: ignore[arg-type]
55
56
57def get_trace_url():
58    # type: () -> str
59    """Return the Agent URL computed from the environment.
60
61    Raises a ``ValueError`` if the URL is not supported by the Agent.
62    """
63    user_supplied_host = get_trace_hostname(None) is not None
64    user_supplied_port = get_trace_port(None) is not None
65
66    url = os.environ.get("DD_TRACE_AGENT_URL")
67
68    if not url:
69        if user_supplied_host or user_supplied_port:
70            url = "http://%s:%s" % (get_trace_hostname(), get_trace_port())
71        elif os.path.exists("/var/run/datadog/apm.socket"):
72            url = "unix://%s" % (DEFAULT_UNIX_TRACE_PATH)
73        else:
74            url = DEFAULT_TRACE_URL
75
76    return url
77
78
79def get_stats_url():
80    # type: () -> str
81    user_supplied_host = get_stats_hostname(None) is not None
82    user_supplied_port = get_stats_port(None) is not None
83
84    url = get_env("dogstatsd", "url", default=None)
85
86    if not url:
87        if user_supplied_host or user_supplied_port:
88            url = "udp://{}:{}".format(get_stats_hostname(), get_stats_port())
89        elif os.path.exists("/var/run/datadog/dsd.socket"):
90            url = "unix://%s" % (DEFAULT_UNIX_DSD_PATH)
91        else:
92            url = "udp://{}:{}".format(get_stats_hostname(), get_stats_port())
93    return url
94
95
96def verify_url(url):
97    # type: (str) -> parse.ParseResult
98    """Verify that a URL can be used to communicate with the Datadog Agent.
99    Returns a parse.ParseResult.
100    Raises a ``ValueError`` if the URL is not supported by the Agent.
101    """
102    parsed = parse.urlparse(url)
103    schemes = ("http", "https", "unix")
104    if parsed.scheme not in schemes:
105        raise ValueError(
106            "Unsupported protocol '%s' in Agent URL '%s'. Must be one of: %s" % (parsed.scheme, url, ", ".join(schemes))
107        )
108    elif parsed.scheme in ["http", "https"] and not parsed.hostname:
109        raise ValueError("Invalid hostname in Agent URL '%s'" % url)
110    elif parsed.scheme == "unix" and not parsed.path:
111        raise ValueError("Invalid file path in Agent URL '%s'" % url)
112
113    return parsed
114
115
116def get_connection(url, timeout=DEFAULT_TIMEOUT):
117    # type: (str, float) -> ConnectionType
118    """Return an HTTP connection to the given URL."""
119    parsed = verify_url(url)
120    hostname = parsed.hostname or ""
121    path = parsed.path or "/"
122
123    if parsed.scheme == "https":
124        return HTTPSConnection.with_base_path(hostname, parsed.port, base_path=path, timeout=timeout)
125    elif parsed.scheme == "http":
126        return HTTPConnection.with_base_path(hostname, parsed.port, base_path=path, timeout=timeout)
127    elif parsed.scheme == "unix":
128        return UDSHTTPConnection(path, hostname, parsed.port, timeout=timeout)
129
130    raise ValueError("Unsupported protocol '%s'" % parsed.scheme)
131