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 utils
32from acloud.internal.lib import ssh as ssh_object
33from acloud.internal.lib.adb_tools import AdbTools
34from acloud.list import list as list_instance
35from acloud.public import config
36from acloud.public import report
37
38
39logger = logging.getLogger(__name__)
40
41_RE_DISPLAY = re.compile(r"([\d]+)x([\d]+)\s.*")
42_VNC_STARTED_PATTERN = "ssvnc vnc://127.0.0.1:%(vnc_port)d"
43_WEBRTC_PORTS_SEARCH = "".join(
44    [utils.PORT_MAPPING % {"local_port":port["local"],
45                           "target_port":port["target"]}
46     for port in utils.WEBRTC_PORTS_MAPPING])
47
48
49def _IsWebrtcEnable(instance, host_user, host_ssh_private_key_path,
50                    extra_args_ssh_tunnel):
51    """Check local/remote instance webRTC is enable.
52
53    Args:
54        instance: Local/Remote Instance object.
55        host_user: String of user login into the instance.
56        host_ssh_private_key_path: String of host key for logging in to the
57                                   host.
58        extra_args_ssh_tunnel: String, extra args for ssh tunnel connection.
59
60    Returns:
61        Boolean: True if cf_runtime_cfg.enable_webrtc is True.
62    """
63    if instance.islocal:
64        return instance.cf_runtime_cfg.enable_webrtc
65    ssh = ssh_object.Ssh(ip=ssh_object.IP(ip=instance.ip), user=host_user,
66                         ssh_private_key_path=host_ssh_private_key_path,
67                         extra_args_ssh_tunnel=extra_args_ssh_tunnel)
68    remote_cuttlefish_config = os.path.join(constants.REMOTE_LOG_FOLDER,
69                                            constants.CUTTLEFISH_CONFIG_FILE)
70    raw_data = ssh.GetCmdOutput("cat " + remote_cuttlefish_config)
71    try:
72        cf_runtime_cfg = cvd_runtime_config.CvdRuntimeConfig(
73            raw_data=raw_data.strip())
74        return cf_runtime_cfg.enable_webrtc
75    except errors.ConfigError:
76        logger.debug("No cuttlefish config[%s] found!",
77                     remote_cuttlefish_config)
78    return False
79
80
81def _WebrtcPortOccupied():
82    """To decide whether need to release port.
83
84    Remote webrtc instance will create a ssh tunnel which may conflict with
85    local webrtc instance default port. Searching process cmd in the pattern
86    of _WEBRTC_PORTS_SEARCH to decide whether to release port.
87
88    Return:
89        True if need to release port.
90    """
91    process_output = utils.CheckOutput(constants.COMMAND_PS)
92    for line in process_output.splitlines():
93        match = re.search(_WEBRTC_PORTS_SEARCH, line)
94        if match:
95            return True
96    return False
97
98
99def StartVnc(vnc_port, display):
100    """Start vnc connect to AVD.
101
102    Confirm whether there is already a connection before VNC connection.
103    If there is a connection, it will not be connected. If not, connect it.
104    Before reconnecting, clear old disconnect ssvnc viewer.
105
106    Args:
107        vnc_port: Integer of vnc port number.
108        display: String, vnc connection resolution. e.g., 1080x720 (240)
109    """
110    vnc_started_pattern = _VNC_STARTED_PATTERN % {"vnc_port": vnc_port}
111    if not utils.IsCommandRunning(vnc_started_pattern):
112        #clean old disconnect ssvnc viewer.
113        utils.CleanupSSVncviewer(vnc_port)
114
115        match = _RE_DISPLAY.match(display)
116        if match:
117            utils.LaunchVncClient(vnc_port, match.group(1), match.group(2))
118        else:
119            utils.LaunchVncClient(vnc_port)
120
121
122def AddPublicSshRsaToInstance(cfg, user, instance_name):
123    """Add the public rsa key to the instance's metadata.
124
125    When the public key doesn't exist in the metadata, it will add it.
126
127    Args:
128        cfg: An AcloudConfig instance.
129        user: String, the ssh username to access instance.
130        instance_name: String, instance name.
131    """
132    credentials = auth.CreateCredentials(cfg)
133    compute_client = android_compute_client.AndroidComputeClient(
134        cfg, credentials)
135    compute_client.AddSshRsaInstanceMetadata(
136        user,
137        cfg.ssh_public_key_path,
138        instance_name)
139
140
141@utils.TimeExecute(function_description="Reconnect instances")
142def ReconnectInstance(ssh_private_key_path,
143                      instance,
144                      reconnect_report,
145                      extra_args_ssh_tunnel=None,
146                      connect_vnc=True):
147    """Reconnect to the specified instance.
148
149    It will:
150     - re-establish ssh tunnels for adb/vnc port forwarding
151     - re-establish adb connection
152     - restart vnc client
153     - update device information in reconnect_report
154
155    Args:
156        ssh_private_key_path: Path to the private key file.
157                              e.g. ~/.ssh/acloud_rsa
158        instance: list.Instance() object.
159        reconnect_report: Report object.
160        extra_args_ssh_tunnel: String, extra args for ssh tunnel connection.
161        connect_vnc: Boolean, True will launch vnc.
162
163    Raises:
164        errors.UnknownAvdType: Unable to reconnect to instance of unknown avd
165                               type.
166    """
167    if instance.avd_type not in utils.AVD_PORT_DICT:
168        raise errors.UnknownAvdType("Unable to reconnect to instance (%s) of "
169                                    "unknown avd type: %s" %
170                                    (instance.name, instance.avd_type))
171
172    adb_cmd = AdbTools(instance.adb_port)
173    vnc_port = instance.vnc_port
174    adb_port = instance.adb_port
175    # ssh tunnel is up but device is disconnected on adb
176    if instance.ssh_tunnel_is_connected and not adb_cmd.IsAdbConnectionAlive():
177        adb_cmd.DisconnectAdb()
178        adb_cmd.ConnectAdb()
179    # ssh tunnel is down and it's a remote instance
180    elif not instance.ssh_tunnel_is_connected and not instance.islocal:
181        adb_cmd.DisconnectAdb()
182        forwarded_ports = utils.AutoConnect(
183            ip_addr=instance.ip,
184            rsa_key_file=ssh_private_key_path,
185            target_vnc_port=utils.AVD_PORT_DICT[instance.avd_type].vnc_port,
186            target_adb_port=utils.AVD_PORT_DICT[instance.avd_type].adb_port,
187            ssh_user=constants.GCE_USER,
188            extra_args_ssh_tunnel=extra_args_ssh_tunnel)
189        vnc_port = forwarded_ports.vnc_port
190        adb_port = forwarded_ports.adb_port
191    if _IsWebrtcEnable(instance,
192                       constants.GCE_USER,
193                       ssh_private_key_path,
194                       extra_args_ssh_tunnel):
195        if instance.islocal:
196            if _WebrtcPortOccupied():
197                raise errors.PortOccupied("\nReconnect to a local webrtc instance "
198                                          "is not work because remote webrtc "
199                                          "instance has established ssh tunnel "
200                                          "which occupied local webrtc instance "
201                                          "port. If you want to connect to a "
202                                          "local-instance of webrtc. please run "
203                                          "'acloud create --local-instance "
204                                          "--autoconnect webrtc' directly.")
205        else:
206            utils.EstablishWebRTCSshTunnel(
207                ip_addr=instance.ip,
208                rsa_key_file=ssh_private_key_path,
209                ssh_user=constants.GCE_USER,
210                extra_args_ssh_tunnel=extra_args_ssh_tunnel)
211        utils.LaunchBrowser(constants.WEBRTC_LOCAL_HOST,
212                            constants.WEBRTC_LOCAL_PORT)
213    elif(vnc_port and connect_vnc):
214        StartVnc(vnc_port, instance.display)
215
216    device_dict = {
217        constants.IP: instance.ip,
218        constants.INSTANCE_NAME: instance.name,
219        constants.VNC_PORT: vnc_port,
220        constants.ADB_PORT: adb_port
221    }
222
223    if vnc_port and adb_port:
224        reconnect_report.AddData(key="devices", value=device_dict)
225    else:
226        # We use 'ps aux' to grep adb/vnc fowarding port from ssh tunnel
227        # command. Therefore we report failure here if no vnc_port and
228        # adb_port found.
229        reconnect_report.AddData(key="device_failing_reconnect", value=device_dict)
230        reconnect_report.AddError(instance.name)
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                          connect_vnc=(args.autoconnect is True))
263
264    utils.PrintDeviceSummary(reconnect_report)
265