1#
2# Copyright (C) 2017 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
16import datetime
17import logging
18import os
19import shutil
20import tempfile
21
22from vts.utils.python.common import cmd_utils
23from vts.utils.python.gcs import gcs_api_utils
24
25
26def NotNoneStr(item):
27    '''Convert a variable to string only if it is not None'''
28    return str(item) if item is not None else None
29
30
31class ReportFileUtil(object):
32    '''Utility class for report file saving.
33
34    Contains methods to save report files or read incremental parts of
35    report files to a destination folder and get URLs.
36    Used by profiling util, systrace util, and host log reporting.
37
38    Attributes:
39        _flatten_source_dir: bool, whether to flatten source directory
40                             structure in destination directory. Current
41                             implementation assumes no duplicated fine names
42        _use_destination_date_dir: bool, whether to create date directory
43                                   in destination directory
44        _source_dir: string, source directory that contains report files
45        _destination_dir: string, the GCS destination bucket name.
46        _url_prefix: string, a prefix added to relative destination file paths.
47                     If set to None, will use parent directory path.
48        _use_gcs: bool, whether or not this ReportFileUtil is using GCS.
49        _gcs_api_utils: GcsApiUtils object used by the ReportFileUtil object.
50        _gcs_available: bool, whether or not the GCS agent is available.
51    '''
52
53    def __init__(self,
54                 flatten_source_dir=False,
55                 use_destination_date_dir=False,
56                 source_dir=None,
57                 destination_dir=None,
58                 url_prefix=None,
59                 gcs_key_path=None):
60        """Initializes the ReportFileUtils object.
61
62        Args:
63            flatten_source_dir: bool, whether or not flatten the directory structure.
64            use_destination_date_dir: bool, whether or not use date as part of name,
65            source_dir: string, path to the source directory.
66            destination_dir: string, path to the destination directory.
67            url_prefix: string, prefix of the url used to upload the link to dashboard.
68            gcs_key_path: string, path to the GCS key file.
69        """
70        source_dir = NotNoneStr(source_dir)
71        destination_dir = NotNoneStr(destination_dir)
72        url_prefix = NotNoneStr(url_prefix)
73
74        self._flatten_source_dir = flatten_source_dir
75        self._use_destination_date_dir = use_destination_date_dir
76        self._source_dir = source_dir
77        self._destination_dir = destination_dir
78        self._url_prefix = url_prefix
79        self._use_gcs = False
80
81        if gcs_key_path is not None:
82            self._use_gcs = True
83            self._gcs_api_utils = gcs_api_utils.GcsApiUtils(
84                gcs_key_path, destination_dir)
85            self._gcs_available = self._gcs_api_utils.Enabled
86
87    def _ConvertReportPath(self,
88                           src_path,
89                           root_dir=None,
90                           new_file_name=None,
91                           file_name_prefix=None):
92        '''Convert report source file path to destination path and url.
93
94        Args:
95            src_path: string, source report file path.
96            new_file_name: string, new file name to use on destination.
97            file_name_prefix: string, prefix added to destination file name.
98                              if new_file_name is set, prefix will be added
99                              to new_file_name as well.
100
101        Returns:
102            tuple(string, string), containing destination path and url
103        '''
104        root_dir = NotNoneStr(root_dir)
105        new_file_name = NotNoneStr(new_file_name)
106        file_name_prefix = NotNoneStr(file_name_prefix)
107
108        dir_path = os.path.dirname(src_path)
109
110        relative_path = os.path.basename(src_path)
111        if new_file_name:
112            relative_path = new_file_name
113        if file_name_prefix:
114            relative_path = file_name_prefix + relative_path
115        if not self._flatten_source_dir and root_dir:
116            relative_path = os.path.join(
117                os.path.relpath(dir_path, root_dir), relative_path)
118        if self._use_destination_date_dir:
119            now = datetime.datetime.now()
120            date = now.strftime('%Y-%m-%d')
121            relative_path = os.path.join(date, relative_path)
122
123        if self._use_gcs:
124            dest_path = relative_path
125        else:
126            dest_path = os.path.join(self._destination_dir, relative_path)
127
128        url = dest_path
129        if self._url_prefix is not None:
130            url = self._url_prefix + relative_path
131        return dest_path, url
132
133    def _PushReportFile(self, src_path, dest_path):
134        '''Push a report file to destination.
135
136        Args:
137            src_path: string, source path of report file
138            dest_path: string, destination path of report file
139        '''
140        logging.info('Uploading log %s to %s.', src_path, dest_path)
141
142        src_path = NotNoneStr(src_path)
143        dest_path = NotNoneStr(dest_path)
144
145        parent_dir = os.path.dirname(dest_path)
146        if not os.path.exists(parent_dir):
147            try:
148                os.makedirs(parent_dir)
149            except OSError as e:
150                logging.exception(e)
151        shutil.copy(src_path, dest_path)
152
153    def _PushReportFileGcs(self, src_path, dest_path):
154        """Upload args src file to the bucket in Google Cloud Storage.
155
156        Args:
157            src_path: string, source path of report file
158            dest_path: string, destination path of report file
159        """
160        if not self._gcs_available:
161            logging.error('Logs not being uploaded.')
162            return
163
164        logging.info('Uploading log %s to %s.', src_path, dest_path)
165
166        src_path = NotNoneStr(src_path)
167        dest_path = NotNoneStr(dest_path)
168
169        # Copy snapshot to temp as GCS will not handle dynamic files.
170        temp_dir = tempfile.mkdtemp()
171        shutil.copy(src_path, temp_dir)
172        src_path = os.path.join(temp_dir, os.path.basename(src_path))
173        logging.debug('Snapshot of logs: %s', src_path)
174
175        try:
176            self._gcs_api_utils.UploadFile(src_path, dest_path)
177        except IOError as e:
178            logging.exception(e)
179        finally:
180            logging.debug('removing temporary directory')
181            try:
182                shutil.rmtree(temp_dir)
183            except OSError as e:
184                logging.exception(e)
185
186    def SaveReport(self, src_path, new_file_name=None, file_name_prefix=None):
187        '''Save report file to destination.
188
189        Args:
190            src_path: string, source report file path.
191            new_file_name: string, new file name to use on destination.
192            file_name_prefix: string, prefix added to destination file name.
193                              if new_file_name is set, prefix will be added
194                              to new_file_name as well.
195
196        Returns:
197            string, destination URL of saved report file.
198            If url_prefix is set to None, will return destination path of
199            report files. If error happens during read or write operation,
200            this method will return None.
201        '''
202        src_path = NotNoneStr(src_path)
203        new_file_name = NotNoneStr(new_file_name)
204        file_name_prefix = NotNoneStr(file_name_prefix)
205
206        try:
207            dest_path, url = self._ConvertReportPath(
208                src_path,
209                new_file_name=new_file_name,
210                file_name_prefix=file_name_prefix)
211            if self._use_gcs:
212                self._PushReportFileGcs(src_path, dest_path)
213            else:
214                self._PushReportFile(src_path, dest_path)
215
216            return url
217        except IOError as e:
218            logging.exception(e)
219
220    def SaveReportsFromDirectory(self,
221                                 source_dir=None,
222                                 file_name_prefix=None,
223                                 file_path_filters=None,
224                                 dryrun=False):
225        '''Save report files from source directory to destination.
226
227        Args:
228            source_dir: string, source directory where report files are stored.
229                        if None, class attribute source_dir will be used.
230                        Default is None.
231            file_name_prefix: string, prefix added to destination file name
232            file_path_filter: function, a functions that return True (pass) or
233                              False (reject) given original file path.
234            dryrun: bool, whether to perform a dry run to get urls only.
235
236        Returns:
237            A list of string, containing destination URLs of saved report files.
238            If url_prefix is set to None, will return destination path of
239            report files. If error happens during read or write operation,
240            this method will return None.
241        '''
242        source_dir = NotNoneStr(source_dir)
243        file_name_prefix = NotNoneStr(file_name_prefix)
244        if not source_dir:
245            source_dir = self._source_dir
246
247        try:
248            urls = []
249
250            for (dirpath, dirnames, filenames) in os.walk(
251                    source_dir, followlinks=False):
252                for filename in filenames:
253                    src_path = os.path.join(dirpath, filename)
254                    dest_path, url = self._ConvertReportPath(
255                        src_path,
256                        root_dir=source_dir,
257                        file_name_prefix=file_name_prefix)
258
259                    if file_path_filters and not file_path_filters(src_path):
260                        continue
261
262                    #TODO(yuexima): handle duplicated destination file names
263                    if not dryrun:
264                        if self._use_gcs:
265                            self._PushReportFileGcs(src_path, dest_path)
266                        else:
267                            self._PushReportFile(src_path, dest_path)
268                    urls.append(url)
269
270            return urls
271        except IOError as e:
272            logging.exception(e)
273