1# Copyright 2018 - The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14r"""Reconnect entry point.
15
16Reconnect will:
17 - re-establish ssh tunnels for adb/vnc port forwarding for a remote instance
18 - adb connect to forwarded ssh port for remote instance
19 - restart vnc for remote/local instances
20"""
21
22import logging
23import os
24import re
25
26from acloud import errors
27from acloud.internal import constants
28from acloud.internal.lib import auth
29from acloud.internal.lib import android_compute_client
30from acloud.internal.lib import cvd_runtime_config
31from acloud.internal.lib import gcompute_client
32from acloud.internal.lib import utils
33from acloud.internal.lib import ssh as ssh_object
34from acloud.internal.lib.adb_tools import AdbTools
35from acloud.list import list as list_instance
36from acloud.public import config
37from acloud.public import report
38
39
40logger = logging.getLogger(__name__)
41
42_RE_DISPLAY = re.compile(r"([\d]+)x([\d]+)\s.*")
43_VNC_STARTED_PATTERN = "ssvnc vnc://127.0.0.1:%(vnc_port)d"
44
45
46def _IsWebrtcEnable(instance, host_user, host_ssh_private_key_path,
47                    extra_args_ssh_tunnel):
48    """Check local/remote instance webRTC is enable.
49
50    Args:
51        instance: Local/Remote Instance object.
52        host_user: String of user login into the instance.
53        host_ssh_private_key_path: String of host key for logging in to the
54                                   host.
55        extra_args_ssh_tunnel: String, extra args for ssh tunnel connection.
56
57    Returns:
58        Boolean: True if cf_runtime_cfg.enable_webrtc is True.
59    """
60    if instance.islocal:
61        return instance.cf_runtime_cfg.enable_webrtc
62    ssh = ssh_object.Ssh(ip=ssh_object.IP(ip=instance.ip), user=host_user,
63                         ssh_private_key_path=host_ssh_private_key_path,
64                         extra_args_ssh_tunnel=extra_args_ssh_tunnel)
65    remote_cuttlefish_config = os.path.join(constants.REMOTE_LOG_FOLDER,
66                                            constants.CUTTLEFISH_CONFIG_FILE)
67    raw_data = ssh.GetCmdOutput("cat " + remote_cuttlefish_config)
68    try:
69        cf_runtime_cfg = cvd_runtime_config.CvdRuntimeConfig(
70            raw_data=raw_data.strip())
71        return cf_runtime_cfg.enable_webrtc
72    except errors.ConfigError:
73        logger.debug("No cuttlefish config[%s] found!",
74                     remote_cuttlefish_config)
75    return False
76
77
78def StartVnc(vnc_port, display):
79    """Start vnc connect to AVD.
80
81    Confirm whether there is already a connection before VNC connection.
82    If there is a connection, it will not be connected. If not, connect it.
83    Before reconnecting, clear old disconnect ssvnc viewer.
84
85    Args:
86        vnc_port: Integer of vnc port number.
87        display: String, vnc connection resolution. e.g., 1080x720 (240)
88    """
89    vnc_started_pattern = _VNC_STARTED_PATTERN % {"vnc_port": vnc_port}
90    if not utils.IsCommandRunning(vnc_started_pattern):
91        #clean old disconnect ssvnc viewer.
92        utils.CleanupSSVncviewer(vnc_port)
93
94        match = _RE_DISPLAY.match(display)
95        if match:
96            utils.LaunchVncClient(vnc_port, match.group(1), match.group(2))
97        else:
98            utils.LaunchVncClient(vnc_port)
99
100
101def AddPublicSshRsaToInstance(cfg, user, instance_name):
102    """Add the public rsa key to the instance's metadata.
103
104    When the public key doesn't exist in the metadata, it will add it.
105
106    Args:
107        cfg: An AcloudConfig instance.
108        user: String, the ssh username to access instance.
109        instance_name: String, instance name.
110    """
111    credentials = auth.CreateCredentials(cfg)
112    compute_client = android_compute_client.AndroidComputeClient(
113        cfg, credentials)
114    compute_client.AddSshRsaInstanceMetadata(
115        user,
116        cfg.ssh_public_key_path,
117        instance_name)
118
119
120@utils.TimeExecute(function_description="Reconnect instances")
121def ReconnectInstance(ssh_private_key_path,
122                      instance,
123                      reconnect_report,
124                      extra_args_ssh_tunnel=None,
125                      autoconnect=None,
126                      connect_hostname=None):
127    """Reconnect to the specified instance.
128
129    It will:
130     - re-establish ssh tunnels for adb/vnc port forwarding
131     - re-establish adb connection
132     - restart vnc client
133     - update device information in reconnect_report
134
135    Args:
136        ssh_private_key_path: Path to the private key file.
137                              e.g. ~/.ssh/acloud_rsa
138        instance: list.Instance() object.
139        reconnect_report: Report object.
140        extra_args_ssh_tunnel: String, extra args for ssh tunnel connection.
141        autoconnect: String, for decide whether to launch vnc/browser or not.
142        connect_hostname: String, the hostname for ssh connect.
143
144    Raises:
145        errors.UnknownAvdType: Unable to reconnect to instance of unknown avd
146                               type.
147    """
148    if instance.avd_type not in utils.AVD_PORT_DICT:
149        raise errors.UnknownAvdType("Unable to reconnect to instance (%s) of "
150                                    "unknown avd type: %s" %
151                                    (instance.name, instance.avd_type))
152    # Ignore extra ssh tunnel to connect with hostname.
153    if connect_hostname:
154        extra_args_ssh_tunnel = None
155
156    adb_cmd = AdbTools(instance.adb_port)
157    vnc_port = instance.vnc_port
158    adb_port = instance.adb_port
159    webrtc_port = instance.webrtc_port
160    # ssh tunnel is up but device is disconnected on adb
161    if instance.ssh_tunnel_is_connected and not adb_cmd.IsAdbConnectionAlive():
162        adb_cmd.DisconnectAdb()
163        adb_cmd.ConnectAdb()
164    # ssh tunnel is down and it's a remote instance
165    elif not instance.ssh_tunnel_is_connected and not instance.islocal:
166        adb_cmd.DisconnectAdb()
167        forwarded_ports = utils.AutoConnect(
168            ip_addr=connect_hostname or instance.ip,
169            rsa_key_file=ssh_private_key_path,
170            target_vnc_port=utils.AVD_PORT_DICT[instance.avd_type].vnc_port,
171            target_adb_port=utils.AVD_PORT_DICT[instance.avd_type].adb_port,
172            ssh_user=constants.GCE_USER,
173            extra_args_ssh_tunnel=extra_args_ssh_tunnel)
174        vnc_port = forwarded_ports.vnc_port
175        adb_port = forwarded_ports.adb_port
176    if autoconnect is constants.INS_KEY_WEBRTC:
177        if not instance.islocal:
178            webrtc_port = utils.GetWebrtcPortFromSSHTunnel(instance.ip)
179            if not webrtc_port:
180                webrtc_port = utils.PickFreePort()
181                utils.EstablishWebRTCSshTunnel(
182                    ip_addr=connect_hostname or instance.ip,
183                    webrtc_local_port=webrtc_port,
184                    rsa_key_file=ssh_private_key_path,
185                    ssh_user=constants.GCE_USER,
186                    extra_args_ssh_tunnel=extra_args_ssh_tunnel)
187        utils.LaunchBrowser(constants.WEBRTC_LOCAL_HOST,
188                            webrtc_port)
189    elif vnc_port and autoconnect is constants.INS_KEY_VNC:
190        StartVnc(vnc_port, instance.display)
191
192    device_dict = {
193        constants.IP: instance.ip,
194        constants.INSTANCE_NAME: instance.name,
195        constants.VNC_PORT: vnc_port,
196        constants.ADB_PORT: adb_port
197    }
198    if adb_port and not instance.islocal:
199        device_dict[constants.DEVICE_SERIAL] = (
200            constants.REMOTE_INSTANCE_ADB_SERIAL % adb_port)
201
202    if (vnc_port or webrtc_port) and adb_port:
203        reconnect_report.AddData(key="devices", value=device_dict)
204    else:
205        # We use 'ps aux' to grep adb/vnc fowarding port from ssh tunnel
206        # command. Therefore we report failure here if no vnc_port and
207        # adb_port found.
208        reconnect_report.AddData(key="device_failing_reconnect", value=device_dict)
209        reconnect_report.AddError(instance.name)
210
211
212def GetSshConnectHostname(cfg, instance):
213    """Get ssh connect hostname.
214
215    Get GCE hostname with specific rule for cloudtop users.
216
217    Args:
218        cfg: AcloudConfig object.
219        instance: list.Instance() object.
220
221    Returns:
222        String of hostname for ssh connect. None is for not connect with
223        hostname such as local instance mode.
224    """
225    if instance.islocal:
226        return None
227    if cfg.connect_hostname:
228        return gcompute_client.GetGCEHostName(
229            cfg.project, instance.name, cfg.zone)
230    return None
231
232
233def Run(args):
234    """Run reconnect.
235
236    Args:
237        args: Namespace object from argparse.parse_args.
238    """
239    cfg = config.GetAcloudConfig(args)
240    instances_to_reconnect = []
241    if args.instance_names is not None:
242        # user input instance name to get instance object.
243        instances_to_reconnect = list_instance.GetInstancesFromInstanceNames(
244            cfg, args.instance_names)
245    if not instances_to_reconnect:
246        instances_to_reconnect = list_instance.ChooseInstances(cfg, args.all)
247
248    reconnect_report = report.Report(command="reconnect")
249    for instance in instances_to_reconnect:
250        if instance.avd_type not in utils.AVD_PORT_DICT:
251            utils.PrintColorString("Skipping reconnect of instance %s due to "
252                                   "unknown avd type (%s)." %
253                                   (instance.name, instance.avd_type),
254                                   utils.TextColors.WARNING)
255            continue
256        if not instance.islocal:
257            AddPublicSshRsaToInstance(cfg, constants.GCE_USER, instance.name)
258        ReconnectInstance(cfg.ssh_private_key_path,
259                          instance,
260                          reconnect_report,
261                          cfg.extra_args_ssh_tunnel,
262                          autoconnect=(args.autoconnect or instance.autoconnect),
263                          connect_hostname=GetSshConnectHostname(cfg, instance))
264
265    utils.PrintDeviceSummary(reconnect_report)
266