1# Copyright 2014 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import logging
6import socket
7
8from autotest_lib.client.common_lib import error
9
10# See server/cros/network/wifi_test_context_manager.py for commandline
11# flags to control IP addresses in WiFi tests.
12DEFAULT_FAILURE_MESSAGE = (
13        'Cannot infer DNS name of companion device from an IP address.')
14ATTENUATOR_FAILURE_MESSAGE = (
15        'Cannot infer DNS name of WiFi variable attenuator from a client IP '
16        'address.  Use --atten_addr=<ip or dns name>')
17ROUTER_FAILURE_MESSAGE = (
18        'Cannot infer DNS name of WiFi router from a client IP address.')
19PCAP_FAILURE_MESSAGE = (
20        'Cannot infer DNS name of Packet Capturer from a client IP address.')
21
22
23def is_ip_address(hostname):
24    """Infers whether |hostname| could be an IP address.
25
26    @param hostname: string DNS name or IP address.
27    @return True iff hostname is a valid IP address.
28
29    """
30    try:
31        socket.inet_aton(hostname)
32        return True
33    except socket.error:
34        return False
35
36
37def get_companion_device_addr(client_hostname,
38                              suffix,
39                              cmdline_override=None,
40                              not_dnsname_msg=DEFAULT_FAILURE_MESSAGE,
41                              allow_failure=False):
42    """Build a usable hostname for a test companion device from the client name.
43
44    Optionally, override the generated name with a commandline provided version.
45
46    @param client_hostname: string DNS name of device under test (the client).
47    @param suffix: string suffix to append to the client hostname.
48    @param cmdline_override: optional DNS name of companion device.  If this is
49            given, it overrides the generated client based hostname.
50    @param not_dnsname_msg: string message to include in the exception raised
51            if the client hostname is found to be an IP address rather than a
52            DNS name.
53    @param allow_failure: boolean True iff we should return None on failure to
54            infer a DNS name.
55    @return string DNS name of companion device or None if |allow_failure|
56            is True and no DNS name can be inferred.
57
58    """
59    if cmdline_override is not None:
60        return cmdline_override
61    if is_ip_address(client_hostname):
62        logging.error('%r looks like an IP address?', client_hostname)
63        if allow_failure:
64            return None
65        raise error.TestError(not_dnsname_msg)
66    parts = client_hostname.split('.', 1)
67    parts[0] = parts[0] + suffix
68    return '.'.join(parts)
69
70
71def get_router_addr(client_hostname, cmdline_override=None):
72    """Build a hostname for a WiFi router from the client hostname.
73
74    Optionally override that hostname with the provided command line hostname.
75
76    @param client_hostname: string DNS name of the client.
77    @param cmdline_override: string DNS name of the router provided
78            via commandline arguments.
79    @return usable DNS name for router host.
80
81    """
82    return get_companion_device_addr(
83            client_hostname,
84            '-router',
85            cmdline_override=cmdline_override,
86            not_dnsname_msg=ROUTER_FAILURE_MESSAGE)
87
88
89def get_pcap_addr(client_hostname,
90                  cmdline_override=None,
91                  allow_failure=False):
92    """Build a hostname for a packet capturer from the client hostname.
93
94    @param client_hostname: string DNS name of the client.
95    @param cmdline_override: string DNS name of the packet capturer provided
96            via commandline arguments.
97    @return usable DNS name for capturer host or None.
98
99    """
100    return get_companion_device_addr(
101            client_hostname,
102            '-pcap',
103            cmdline_override=cmdline_override,
104            not_dnsname_msg=PCAP_FAILURE_MESSAGE,
105            allow_failure=allow_failure)
106
107
108def get_attenuator_addr(client_hostname,
109                        cmdline_override=None,
110                        allow_failure=False):
111    """Build a hostname for a WiFi variable attenuator from the client hostname.
112
113    Optionally override that hostname with the provided command line hostname.
114
115    @param client_hostname: string DNS name of the client.
116    @param cmdline_override: string DNS name of the variable attenuator
117            controller provided via commandline arguments.
118    @param allow_failure: boolean True iff we should return None on failure to
119            infer a DNS name.
120    @return usable DNS name for attenuator controller.
121
122    """
123    return get_companion_device_addr(
124            client_hostname,
125            '-attenuator',
126            cmdline_override=cmdline_override,
127            not_dnsname_msg=ATTENUATOR_FAILURE_MESSAGE,
128            allow_failure=allow_failure)
129