1#
2# Copyright (C) 2018 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16
17import getpass
18import json
19import logging
20import os
21from os.path import expanduser
22import re
23import shutil
24import tempfile
25
26from acloud.public import acloud_main
27from host_controller.acloud import acloud_config
28from vts.utils.python.common import cmd_utils
29
30DEFAULT_BRANCH = 'git_master'
31DEFAULT_BUILD_TARGET = 'gce_x86_phone-userdebug_fastbuild3c_linux'
32
33#TODO(yuexima): add full support to multiple instances per host
34
35
36class ACloudClient(object):
37    '''Helper class to manage access to the acloud module.'''
38
39    def __init__(self):
40        tmpdir_base = os.path.join(os.getcwd(), "tmp")
41        if not os.path.exists(tmpdir_base):
42            os.mkdir(tmpdir_base)
43        self._tmpdir = tempfile.mkdtemp(dir=tmpdir_base)
44
45    def __del__(self):
46        """Deletes the temp dir if still set."""
47        if self._tmpdir:
48            shutil.rmtree(self._tmp_dirpath)
49            self._tmpdir = None
50
51    def GetCreateCmd(self,
52                     build_id,
53                     branch=None,
54                     build_target=None,
55                     num=1):
56        '''Get acould create command with given build id.
57
58        Args:
59            build_id: string, build_id.
60            branch: string, build branch
61            build_target: string, build target
62            num: int, number of instances to build
63
64        Returns:
65            string, acloud create command.
66        '''
67        if not branch:
68            branch = DEFAULT_BRANCH
69
70        if not build_target:
71            build_target = DEFAULT_BUILD_TARGET
72
73        #TODO latest build id (in the caller class of this function).
74        #TODO use unique log and tmp file location
75        cmd = ('create '
76            '--branch {branch} '
77            '--build_target {build_target} '
78            '--build_id {build_id} '
79            '--config_file {config_file} '
80            '--report_file {report_file} '
81            '--log_file {log_file} '
82            '--email {email} '
83            '--num {num}').format(
84            branch = branch,
85            build_target = build_target,
86            build_id = build_id,
87            config_file = os.path.join(self._tmpdir, 'acloud.config'),
88            report_file = os.path.join(self._tmpdir, 'acloud_report.json'),
89            log_file = os.path.join(self._tmpdir, 'acloud.log'),
90            #TODO use host email address
91            email = getpass.getuser() + '@google.com',
92            num = num
93        )
94        return cmd
95
96    def GetDeleteCmd(self, instance_names):
97        '''Get Acould delete command with given instance names.
98
99        Args:
100            instance_names: list of string, instance names.
101
102        Returns:
103            string, acloud create command.
104        '''
105        cmd = ('delete '
106            '--instance_names {instance_names} '
107            '--config_file {config_file} '
108            '--report_file {report_file} '
109            '--log_file {log_file} '
110            '--email {email}').format(
111            instance_names = ' '.join(instance_names),
112            config_file = os.path.join(self._tmpdir, 'acloud.config'),
113            report_file = os.path.join(self._tmpdir, 'acloud_report.json'),
114            log_file = os.path.join(self._tmpdir, 'acloud.log'),
115            email = getpass.getuser() + '@google.com'
116        )
117        return cmd
118
119    def CreateInstance(self, build_id):
120        '''Create Acould instance with given build id.
121
122        Args:
123            build_id: string, build_id.
124        '''
125        cmd = self.GetCreateCmd(build_id)
126        acloud_main.main(cmd.split())
127
128        report_file = os.path.join(self._tmpdir, 'acloud_report.json')
129        with open(report_file, 'r') as f:
130            report = json.load(f)
131
132        return report['status']=='SUCCESS'
133
134    def PrepareConfig(self, file_path):
135        '''Prepare acloud Acloud config file.
136
137        Args:
138            file_path: string, path to acloud config file.
139        '''
140        config = acloud_config.ACloudConfig()
141        config.Load(file_path)
142        if not config.Validate():
143            logging.error('Failed to prepare acloud config.')
144            return
145
146        config.Save(os.path.join(self._tmpdir, 'acloud.config'))
147
148    def GetInstanceIP(self):
149        '''Get an Acloud instance ip'''
150        #TODO support ip addresses when num > 1 (json parser)
151        report_file = os.path.join(self._tmpdir, 'acloud_report.json')
152
153        with open(report_file, 'r') as f:
154            report = json.load(f)
155
156        return report['data']['devices'][0]['ip']
157
158    def GetInstanceName(self):
159        '''Get an Acloud instance name'''
160        report_file = os.path.join(self._tmpdir, 'acloud_report.json')
161
162        with open(report_file, 'r') as f:
163            report = json.load(f)
164
165        return report['data']['devices'][0]['instance_name']
166
167    def ConnectInstanceToAdb(self, ip=None):
168        '''Connect an Acloud instance to adb
169
170        Args:
171            ip: string, ip address
172        '''
173        if not ip:
174            ip = self.GetInstanceIP()
175
176        cmds = [('ssh -i ~/.ssh/acloud_rsa -o UserKnownHostsFile=/dev/null '
177                 '-o StrictHostKeyChecking=no -L 40000:127.0.0.1:6444 -L '
178                 '30000:127.0.0.1:5555 -N -f -l root %s') % ip,
179                 'adb connect 127.0.0.1:30000']
180
181        for cmd in cmds:
182            print cmd
183            results = cmd_utils.ExecuteShellCommand(cmd)
184            if any(results[cmd_utils.EXIT_CODE]):
185                logging.error("Fail to execute command: %s\nResult:%s" % (cmd,
186                                                                        results))
187                return
188            print 'Acloud instance created and connected to local port 127.0.0.1:30000'
189
190    def DeleteInstance(self, instance_names=None):
191        '''Delete an Acloud instance
192
193        Args:
194            instance_names: string, instance name
195        '''
196        cmd = self.GetDeleteCmd(instance_names)
197        acloud_main.main(cmd.split())
198
199        report_file = os.path.join(self._tmpdir, 'acloud_report.json')
200        with open(report_file, 'r') as f:
201            report = json.load(f)
202
203        return report['status']=='SUCCESS'
204