1#!/usr/bin/env python3
2#
3#  Copyright (C) 2022 The Android Open Source Project
4#
5#  Licensed under the Apache License, Version 2.0 (the "License");
6#  you may not use this file except in compliance with the License.
7#  You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#  Unless required by applicable law or agreed to in writing, software
12#  distributed under the License is distributed on an "AS IS" BASIS,
13#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#  See the License for the specific language governing permissions and
15#  limitations under the License.
16
17"""The script to invoke content_uploader binary to upload artifacts to CAS."""
18import glob
19import logging
20import os
21import subprocess
22import sys
23
24CONTENT_PLOADER_PREBUILT_PATH = 'tools/tradefederation/prebuilts/'
25CONTENT_UPLOADER_BIN = 'content_uploader'
26CONTENT_UPLOADER_TIMEOUT_SECS = 1800 # 30 minutes
27LOG_PATH = 'logs/cas_uploader.log'
28
29def _get_prebuilt_uploader() -> str:
30    uploader = glob.glob(CONTENT_PLOADER_PREBUILT_PATH + '**/' + CONTENT_UPLOADER_BIN,
31                         recursive=True)
32    if not uploader:
33        logging.error('%s not found in Tradefed prebuilt', CONTENT_UPLOADER_BIN)
34        raise ValueError(f'Error: {CONTENT_UPLOADER_BIN} not found in Tradefed prebuilt')
35    return uploader[0]
36
37def _get_env_var(key: str, default=None, check=False):
38    value = os.environ.get(key, default)
39    if check and not value:
40        logging.error('the environment variable %s is not set', key)
41        raise ValueError(f'Error: the environment variable {key} is not set')
42    return value
43
44def _setup_logging() -> str:
45    dist_dir = _get_env_var('DIST_DIR', check=True)
46    log_file = os.path.join(dist_dir, LOG_PATH)
47    logging.basicConfig(
48        level=logging.DEBUG,
49        format='%(asctime)s %(levelname)s %(message)s',
50        filename=log_file,
51    )
52    return log_file
53
54def main():
55    """Call content_uploader and pass all arguments as-is."""
56    log_file = _setup_logging()
57
58    uploader = _get_prebuilt_uploader()
59    arguments = sys.argv[1:]
60
61    try:
62        result = subprocess.run([uploader] + arguments, capture_output=True, text=True,
63                                check=True, timeout=CONTENT_UPLOADER_TIMEOUT_SECS)
64        print(result.stdout)
65    except FileNotFoundError:
66        print(f'content_uploader.py will export logs to: {log_file}')
67        logging.error('Uploader not found: %s', log_file)
68
69if __name__ == '__main__':
70    main()
71