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 logging
18
19from vts.utils.python.common import cmd_utils
20
21
22def GetGsutilPath():
23    """Finds gsutil in PATH.
24
25    Instead of a Python library, gsutil binary is used to avoid packaging GCS
26    PIP package as part of VTS HC (Host Controller).
27
28    Returns:
29        The gsutil file path if found; None otherwise.
30    """
31    sh_stdout, sh_stderr, ret_code = cmd_utils.ExecuteOneShellCommand(
32        "which gsutil")
33    if ret_code == 0:
34        return sh_stdout.strip()
35    else:
36        logging.fatal("`gsutil` doesn't exist on the host; "
37                      "please install Google Cloud SDK before retrying.")
38        return None
39
40
41def IsGcsFile(gsutil_path, url):
42    """Checks whether a given path is for a GCS file.
43
44    Args:
45        gsutil_path: string, the path of a gsutil binary.
46        url: string, the GCS URL. e.g., gs://<bucket>/<file>.
47
48    Returns:
49        True if url is a file, False otherwise.
50    """
51    check_command = "%s stat %s" % (gsutil_path, url)
52    _, _, ret_code = cmd_utils.ExecuteOneShellCommand(check_command)
53    return ret_code == 0
54
55
56def Copy(gsutil_path, src_urls, dst_url):
57    """Copies files between local file system and GCS.
58
59    Args:
60        gsutil_path: string, the path of a gsutil binary.
61        src_urls: string, the source paths or GCS URLs separated by spaces.
62        dst_url: string, the destination path or GCS URL.
63
64    Returns:
65        True if the command succeeded, False otherwise.
66    """
67    copy_command = "%s cp %s %s" % (gsutil_path, src_urls, dst_url)
68    _, _, ret_code = cmd_utils.ExecuteOneShellCommand(copy_command)
69    return ret_code == 0
70
71
72def List(gsutil_path, url):
73    """Lists a directory or file on GCS.
74
75    Args:
76        gsutil_path: string, the path of a gsutil binary.
77        url: string, the GCS URL of the directory or file.
78
79    Returns:
80        list of strings, the GCS URLs of the listed files.
81    """
82    ls_command = "%s ls %s" % (gsutil_path, url)
83    stdout, _, ret_code = cmd_utils.ExecuteOneShellCommand(ls_command)
84    return stdout.strip("\n").split("\n") if ret_code == 0 else []
85
86
87def Remove(gsutil_path, url, recursive=False):
88    """Removes a directory or file on GCS.
89
90    Args:
91        gsutil_path: string, the path of a gsutil binary.
92        url: string, the GCS URL of the directory or file.
93        recursive: boolean, whether to remove the directory recursively.
94
95    Returns:
96        True if the command succeeded, False otherwise.
97    """
98    if "/" not in url.lstrip("gs://").rstrip("/"):
99        logging.error("Cannot remove bucket %s", url)
100        return False
101    rm_command = "%s rm -f%s %s" % (
102        gsutil_path, ("r" if recursive else ""), url)
103    ret_code = cmd_utils.ExecuteOneShellCommand(rm_command)
104    return ret_code == 0
105