• 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
22from __future__ import print_function
23
24import getpass
25import re
26
27from acloud.delete import delete
28from acloud.internal.lib import auth
29from acloud.internal.lib import android_compute_client
30from acloud.internal.lib import utils
31from acloud.internal.lib.adb_tools import AdbTools
32from acloud.list import list as list_instance
33from acloud.public import config
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        delete.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        cfg.zone,
77        user,
78        cfg.ssh_public_key_path,
79        instance_name)
80
81
82def ReconnectInstance(ssh_private_key_path, instance):
83    """Reconnect adb/vnc/ssh to the specified instance.
84
85    Args:
86        ssh_private_key_path: Path to the private key file.
87                              e.g. ~/.ssh/acloud_rsa
88        instance: list.Instance() object.
89    """
90    adb_cmd = AdbTools(instance.forwarding_adb_port)
91    vnc_port = instance.forwarding_vnc_port
92    # ssh tunnel is up but device is disconnected on adb
93    if instance.ssh_tunnel_is_connected and not adb_cmd.IsAdbConnectionAlive():
94        adb_cmd.DisconnectAdb()
95        adb_cmd.ConnectAdb()
96    # ssh tunnel is down and it's a remote instance
97    elif not instance.ssh_tunnel_is_connected and not instance.islocal:
98        adb_cmd.DisconnectAdb()
99        forwarded_ports = utils.AutoConnect(
100            instance.ip,
101            ssh_private_key_path,
102            utils.AVD_PORT_DICT[instance.avd_type].vnc_port,
103            utils.AVD_PORT_DICT[instance.avd_type].adb_port,
104            getpass.getuser())
105        vnc_port = forwarded_ports.vnc_port
106
107    if vnc_port:
108        StartVnc(vnc_port, instance.display)
109
110
111def Run(args):
112    """Run reconnect.
113
114    Args:
115        args: Namespace object from argparse.parse_args.
116    """
117    cfg = config.GetAcloudConfig(args)
118    instances_to_reconnect = []
119    if args.instance_names is not None:
120        # user input instance name to get instance object.
121        instances_to_reconnect = list_instance.GetInstancesFromInstanceNames(
122            cfg, args.instance_names)
123    if not instances_to_reconnect:
124        instances_to_reconnect = list_instance.ChooseInstances(cfg, args.all)
125    for instance in instances_to_reconnect:
126        AddPublicSshRsaToInstance(cfg, getpass.getuser(), instance.name)
127        ReconnectInstance(cfg.ssh_private_key_path, instance)
128