1#!/usr/bin/python -u
2
3"""
4Utility to upload or remove the packages from the packages repository.
5"""
6
7import logging, optparse, os, shutil, sys, tempfile
8import common
9from autotest_lib.client.common_lib import utils as client_utils
10from autotest_lib.client.common_lib import global_config, error
11from autotest_lib.client.common_lib import base_packages, packages
12from autotest_lib.server import utils as server_utils
13
14c = global_config.global_config
15logging.basicConfig(level=logging.DEBUG)
16
17def get_exclude_string(client_dir):
18    '''
19    Get the exclude string for the tar command to exclude specific
20    subdirectories inside client_dir.
21    For profilers we need to exclude everything except the __init__.py
22    file so that the profilers can be imported.
23    '''
24    exclude_string = ('--exclude=deps/* --exclude=tests/* '
25                      '--exclude=site_tests/* --exclude=**.pyc')
26
27    # Get the profilers directory
28    prof_dir = os.path.join(client_dir, 'profilers')
29
30    # Include the __init__.py file for the profilers and exclude all its
31    # subdirectories
32    for f in os.listdir(prof_dir):
33        if os.path.isdir(os.path.join(prof_dir, f)):
34            exclude_string += ' --exclude=profilers/%s' % f
35
36    # The '.' here is needed to zip the files in the current
37    # directory. We use '-C' for tar to change to the required
38    # directory i.e. src_dir and then zip up the files in that
39    # directory(which is '.') excluding the ones in the exclude_dirs
40    exclude_string += " ."
41
42    # TODO(milleral): This is sad and ugly.  http://crbug.com/258161
43    # Surprisingly, |exclude_string| actually means argument list, and
44    # we'd like to package up the current global_config.ini also, so let's
45    # just tack it on here.
46    # Also note that this only works because tar prevents us from un-tarring
47    # files into parent directories.
48    exclude_string += " ../global_config.ini"
49
50    return exclude_string
51
52
53def parse_args():
54    parser = optparse.OptionParser()
55    parser.add_option("-d", "--dependency", help="package the dependency"
56                      " from client/deps directory and upload to the repo",
57                      dest="dep")
58    parser.add_option("-p", "--profiler", help="package the profiler "
59                      "from client/profilers directory and upload to the repo",
60                      dest="prof")
61    parser.add_option("-t", "--test", help="package the test from client/tests"
62                      " or client/site_tests and upload to the repo.",
63                      dest="test")
64    parser.add_option("-c", "--client", help="package the client "
65                      "directory alone without the tests, deps and profilers",
66                      dest="client", action="store_true", default=False)
67    parser.add_option("-f", "--file", help="simply uploads the specified"
68                      "file on to the repo", dest="file")
69    parser.add_option("-r", "--repository", help="the URL of the packages"
70                      "repository location to upload the packages to.",
71                      dest="repo", default=None)
72    parser.add_option("--all", help="Upload all the files locally "
73                      "to all the repos specified in global_config.ini. "
74                      "(includes the client, tests, deps and profilers)",
75                      dest="all", action="store_true", default=False)
76
77    options, args = parser.parse_args()
78    return options, args
79
80
81# Method to upload or remove package depending on the flag passed to it.
82def process_packages(pkgmgr, pkg_type, pkg_names, src_dir,
83                    remove=False):
84    exclude_string = ' .'
85    names = [p.strip() for p in pkg_names.split(',')]
86    for name in names:
87        print "Processing %s ... " % name
88        if pkg_type == 'client':
89            pkg_dir = src_dir
90            exclude_string = get_exclude_string(pkg_dir)
91        elif pkg_type == 'test':
92            # if the package is a test then look whether it is in client/tests
93            # or client/site_tests
94            pkg_dir = os.path.join(get_test_dir(name, src_dir), name)
95        else:
96            # for the profilers and deps
97            pkg_dir = os.path.join(src_dir, name)
98
99        pkg_name = pkgmgr.get_tarball_name(name, pkg_type)
100        if not remove:
101            # Tar the source and upload
102            temp_dir = tempfile.mkdtemp()
103            try:
104                try:
105                    base_packages.check_diskspace(temp_dir)
106                except error.RepoDiskFullError, e:
107                    msg = ("Temporary directory for packages %s does not have "
108                           "enough space available: %s" % (temp_dir, e))
109                    raise error.RepoDiskFullError(msg)
110                tarball_path = pkgmgr.tar_package(pkg_name, pkg_dir,
111                                                  temp_dir, exclude_string)
112                pkgmgr.upload_pkg(tarball_path, update_checksum=True)
113            finally:
114                # remove the temporary directory
115                shutil.rmtree(temp_dir)
116        else:
117            pkgmgr.remove_pkg(pkg_name, remove_checksum=True)
118        print "Done."
119
120
121def tar_packages(pkgmgr, pkg_type, pkg_names, src_dir, temp_dir):
122    """Tar all packages up and return a list of each tar created"""
123    tarballs = []
124    exclude_string = ' .'
125    names = [p.strip() for p in pkg_names.split(',')]
126    for name in names:
127        print "Processing %s ... " % name
128        if pkg_type == 'client':
129            pkg_dir = src_dir
130            exclude_string = get_exclude_string(pkg_dir)
131        elif pkg_type == 'test':
132            # if the package is a test then look whether it is in client/tests
133            # or client/site_tests
134            pkg_dir = os.path.join(get_test_dir(name, src_dir), name)
135        else:
136            # for the profilers and deps
137            pkg_dir = os.path.join(src_dir, name)
138
139        pkg_name = pkgmgr.get_tarball_name(name, pkg_type)
140        tarball_path = pkgmgr.tar_package(pkg_name, pkg_dir,
141                                              temp_dir, exclude_string)
142
143        tarballs.append(tarball_path)
144    return tarballs
145
146
147def process_all_packages(pkgmgr, client_dir, remove=False):
148    """Process a full upload of packages as a directory upload."""
149    dep_dir = os.path.join(client_dir, "deps")
150    prof_dir = os.path.join(client_dir, "profilers")
151    # Directory where all are kept
152    temp_dir = tempfile.mkdtemp()
153    try:
154        base_packages.check_diskspace(temp_dir)
155    except error.RepoDiskFullError, e:
156        print ("Temp destination for packages is full %s, aborting upload: %s"
157               % (temp_dir, e))
158        os.rmdir(temp_dir)
159        sys.exit(1)
160
161    # process tests
162    tests_list = get_subdir_list('tests', client_dir)
163    tests = ','.join(tests_list)
164
165    # process site_tests
166    site_tests_list = get_subdir_list('site_tests', client_dir)
167    site_tests = ','.join(site_tests_list)
168
169    # process deps
170    deps_list = get_subdir_list('deps', client_dir)
171    deps = ','.join(deps_list)
172
173    # process profilers
174    profilers_list = get_subdir_list('profilers', client_dir)
175    profilers = ','.join(profilers_list)
176
177    # Update md5sum
178    if not remove:
179        all_packages = []
180        all_packages.extend(tar_packages(pkgmgr, 'profiler', profilers,
181                                         prof_dir, temp_dir))
182        all_packages.extend(tar_packages(pkgmgr, 'dep', deps, dep_dir,
183                                         temp_dir))
184        all_packages.extend(tar_packages(pkgmgr, 'test', site_tests,
185                                         client_dir, temp_dir))
186        all_packages.extend(tar_packages(pkgmgr, 'test', tests, client_dir,
187                                         temp_dir))
188        all_packages.extend(tar_packages(pkgmgr, 'client', 'autotest',
189                                         client_dir, temp_dir))
190        for package in all_packages:
191            pkgmgr.upload_pkg(package, update_checksum=True)
192        client_utils.run('rm -rf ' + temp_dir)
193    else:
194        process_packages(pkgmgr, 'test', tests, client_dir, remove=remove)
195        process_packages(pkgmgr, 'test', site_tests, client_dir, remove=remove)
196        process_packages(pkgmgr, 'client', 'autotest', client_dir,
197                         remove=remove)
198        process_packages(pkgmgr, 'dep', deps, dep_dir, remove=remove)
199        process_packages(pkgmgr, 'profiler', profilers, prof_dir,
200                         remove=remove)
201
202
203# Get the list of sub directories present in a directory
204def get_subdir_list(name, client_dir):
205    dir_name = os.path.join(client_dir, name)
206    return [f for f in
207            os.listdir(dir_name)
208            if os.path.isdir(os.path.join(dir_name, f)) ]
209
210
211# Look whether the test is present in client/tests and client/site_tests dirs
212def get_test_dir(name, client_dir):
213    names_test = os.listdir(os.path.join(client_dir, 'tests'))
214    names_site_test = os.listdir(os.path.join(client_dir, 'site_tests'))
215    if name in names_test:
216        src_dir = os.path.join(client_dir, 'tests')
217    elif name in names_site_test:
218        src_dir = os.path.join(client_dir, 'site_tests')
219    else:
220        print "Test %s not found" % name
221        sys.exit(0)
222    return src_dir
223
224
225def main():
226    # get options and args
227    options, args = parse_args()
228
229    server_dir = server_utils.get_server_dir()
230    autotest_dir = os.path.abspath(os.path.join(server_dir, '..'))
231
232    # extract the pkg locations from global config
233    repo_urls = c.get_config_value('PACKAGES', 'fetch_location',
234                                   type=list, default=[])
235    upload_paths = c.get_config_value('PACKAGES', 'upload_location',
236                                      type=list, default=[])
237
238    if options.repo:
239        upload_paths.append(options.repo)
240
241    # Having no upload paths basically means you're not using packaging.
242    if not upload_paths:
243        print("No upload locations found. Please set upload_location under"
244              " PACKAGES in the global_config.ini or provide a location using"
245              " the --repository option.")
246        return
247
248    client_dir = os.path.join(autotest_dir, "client")
249
250    # Bail out if the client directory does not exist
251    if not os.path.exists(client_dir):
252        sys.exit(0)
253
254    dep_dir = os.path.join(client_dir, "deps")
255    prof_dir = os.path.join(client_dir, "profilers")
256
257    if len(args) == 0 or args[0] not in ['upload', 'remove']:
258        print("Either 'upload' or 'remove' needs to be specified "
259              "for the package")
260        sys.exit(0)
261
262    if args[0] == 'upload':
263        remove_flag = False
264    elif args[0] == 'remove':
265        remove_flag = True
266    else:
267        # we should not be getting here
268        assert(False)
269
270    pkgmgr = packages.PackageManager(autotest_dir, repo_urls=repo_urls,
271                                     upload_paths=upload_paths,
272                                     run_function_dargs={'timeout':600})
273
274    if options.all:
275        process_all_packages(pkgmgr, client_dir, remove=remove_flag)
276
277    if options.client:
278        process_packages(pkgmgr, 'client', 'autotest', client_dir,
279                         remove=remove_flag)
280
281    if options.dep:
282        process_packages(pkgmgr, 'dep', options.dep, dep_dir,
283                         remove=remove_flag)
284
285    if options.test:
286        process_packages(pkgmgr, 'test', options.test, client_dir,
287                         remove=remove_flag)
288
289    if options.prof:
290        process_packages(pkgmgr, 'profiler', options.prof, prof_dir,
291                         remove=remove_flag)
292
293    if options.file:
294        if remove_flag:
295            pkgmgr.remove_pkg(options.file, remove_checksum=True)
296        else:
297            pkgmgr.upload_pkg(options.file, update_checksum=True)
298
299
300if __name__ == "__main__":
301    main()
302