• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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 re
23
24from acloud import errors
25from acloud.internal import constants
26from acloud.internal.lib import auth
27from acloud.internal.lib import android_compute_client
28from acloud.internal.lib import utils
29from acloud.internal.lib.adb_tools import AdbTools
30from acloud.list import list as list_instance
31from acloud.public import config
32from acloud.public import report
33
34
35_RE_DISPLAY = re.compile(r"([\d]+)x([\d]+)\s.*")
36_VNC_STARTED_PATTERN = "ssvnc vnc://127.0.0.1:%(vnc_port)d"
37
38
39def StartVnc(vnc_port, display):
40    """Start vnc connect to AVD.
41
42    Confirm whether there is already a connection before VNC connection.
43    If there is a connection, it will not be connected. If not, connect it.
44    Before reconnecting, clear old disconnect ssvnc viewer.
45
46    Args:
47        vnc_port: Integer of vnc port number.
48        display: String, vnc connection resolution. e.g., 1080x720 (240)
49    """
50    vnc_started_pattern = _VNC_STARTED_PATTERN % {"vnc_port": vnc_port}
51    if not utils.IsCommandRunning(vnc_started_pattern):
52        #clean old disconnect ssvnc viewer.
53        utils.CleanupSSVncviewer(vnc_port)
54
55        match = _RE_DISPLAY.match(display)
56        if match:
57            utils.LaunchVncClient(vnc_port, match.group(1), match.group(2))
58        else:
59            utils.LaunchVncClient(vnc_port)
60
61
62def AddPublicSshRsaToInstance(cfg, user, instance_name):
63    """Add the public rsa key to the instance's metadata.
64
65    When the public key doesn't exist in the metadata, it will add it.
66
67    Args:
68        cfg: An AcloudConfig instance.
69        user: String, the ssh username to access instance.
70        instance_name: String, instance name.
71    """
72    credentials = auth.CreateCredentials(cfg)
73    compute_client = android_compute_client.AndroidComputeClient(
74        cfg, credentials)
75    compute_client.AddSshRsaInstanceMetadata(
76        user,
77        cfg.ssh_public_key_path,
78        instance_name)
79
80
81@utils.TimeExecute(function_description="Reconnect instances")
82def ReconnectInstance(ssh_private_key_path,
83                      instance,
84                      reconnect_report,
85                      extra_args_ssh_tunnel=None,
86                      connect_vnc=True):
87    """Reconnect to the specified instance.
88
89    It will:
90     - re-establish ssh tunnels for adb/vnc port forwarding
91     - re-establish adb connection
92     - restart vnc client
93     - update device information in reconnect_report
94
95    Args:
96        ssh_private_key_path: Path to the private key file.
97                              e.g. ~/.ssh/acloud_rsa
98        instance: list.Instance() object.
99        reconnect_report: Report object.
100        extra_args_ssh_tunnel: String, extra args for ssh tunnel connection.
101        connect_vnc: Boolean, True will launch vnc.
102
103    Raises:
104        errors.UnknownAvdType: Unable to reconnect to instance of unknown avd
105                               type.
106    """
107    if instance.avd_type not in utils.AVD_PORT_DICT:
108        raise errors.UnknownAvdType("Unable to reconnect to instance (%s) of "
109                                    "unknown avd type: %s" %
110                                    (instance.name, instance.avd_type))
111
112    adb_cmd = AdbTools(instance.adb_port)
113    vnc_port = instance.vnc_port
114    adb_port = instance.adb_port
115    # ssh tunnel is up but device is disconnected on adb
116    if instance.ssh_tunnel_is_connected and not adb_cmd.IsAdbConnectionAlive():
117        adb_cmd.DisconnectAdb()
118        adb_cmd.ConnectAdb()
119    # ssh tunnel is down and it's a remote instance
120    elif not instance.ssh_tunnel_is_connected and not instance.islocal:
121        adb_cmd.DisconnectAdb()
122        forwarded_ports = utils.AutoConnect(
123            ip_addr=instance.ip,
124            rsa_key_file=ssh_private_key_path,
125            target_vnc_port=utils.AVD_PORT_DICT[instance.avd_type].vnc_port,
126            target_adb_port=utils.AVD_PORT_DICT[instance.avd_type].adb_port,
127            ssh_user=constants.GCE_USER,
128            extra_args_ssh_tunnel=extra_args_ssh_tunnel)
129        vnc_port = forwarded_ports.vnc_port
130        adb_port = forwarded_ports.adb_port
131
132    if vnc_port and connect_vnc:
133        StartVnc(vnc_port, instance.display)
134
135    device_dict = {
136        constants.IP: instance.ip,
137        constants.INSTANCE_NAME: instance.name,
138        constants.VNC_PORT: vnc_port,
139        constants.ADB_PORT: adb_port
140    }
141
142    if vnc_port and adb_port:
143        reconnect_report.AddData(key="devices", value=device_dict)
144    else:
145        # We use 'ps aux' to grep adb/vnc fowarding port from ssh tunnel
146        # command. Therefore we report failure here if no vnc_port and
147        # adb_port found.
148        reconnect_report.AddData(key="device_failing_reconnect", value=device_dict)
149        reconnect_report.AddError(instance.name)
150
151
152def Run(args):
153    """Run reconnect.
154
155    Args:
156        args: Namespace object from argparse.parse_args.
157    """
158    cfg = config.GetAcloudConfig(args)
159    instances_to_reconnect = []
160    if args.instance_names is not None:
161        # user input instance name to get instance object.
162        instances_to_reconnect = list_instance.GetInstancesFromInstanceNames(
163            cfg, args.instance_names)
164    if not instances_to_reconnect:
165        instances_to_reconnect = list_instance.ChooseInstances(cfg, args.all)
166
167    reconnect_report = report.Report(command="reconnect")
168    for instance in instances_to_reconnect:
169        if instance.avd_type not in utils.AVD_PORT_DICT:
170            utils.PrintColorString("Skipping reconnect of instance %s due to "
171                                   "unknown avd type (%s)." %
172                                   (instance.name, instance.avd_type),
173                                   utils.TextColors.WARNING)
174            continue
175        if not instance.islocal:
176            AddPublicSshRsaToInstance(cfg, constants.GCE_USER, instance.name)
177        ReconnectInstance(cfg.ssh_private_key_path,
178                          instance,
179                          reconnect_report,
180                          cfg.extra_args_ssh_tunnel,
181                          connect_vnc=(args.autoconnect is True))
182
183    utils.PrintDeviceSummary(reconnect_report)
184