1#!/usr/bin/env python3
2#
3# Copyright (C) 2021 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"""A script to copy outputs from Bazel rules to a user specified dist directory.
17
18This script is only meant to be executed with `bazel run`. `bazel build <this
19script>` doesn't actually copy the files, you'd have to `bazel run` a
20copy_to_dist_dir target.
21
22This script copies files from Bazel's output tree into a directory specified by
23the user. It does not check if the dist dir already contains the file, and will
24simply overwrite it.
25
26One approach is to wipe the dist dir every time this script runs, but that may
27be overly destructive and best left to an explicit rm -rf call outside of this
28script.
29
30Another approach is to error out if the file being copied already exist in the
31dist dir, or perform some kind of content hash checking.
32"""
33
34import argparse
35import collections
36import fnmatch
37import glob
38import logging
39import os
40import pathlib
41import shutil
42import sys
43import tarfile
44
45
46def copy_with_modes(src, dst, mode_overrides):
47    mode_override = None
48    for (pattern, mode) in mode_overrides:
49        if fnmatch.fnmatch(src, pattern):
50            mode_override = mode
51            break
52
53    # Remove destination file that may be write-protected
54    pathlib.Path(dst).unlink(missing_ok=True)
55
56    # Copy the file with copy2 to preserve whatever permissions are set on src
57    shutil.copy2(os.path.abspath(src), dst, follow_symlinks=True)
58
59    if mode_override:
60        os.chmod(dst, mode_override)
61
62
63def ensure_unique_filenames(files):
64    basename_to_srcs_map = collections.defaultdict(list)
65    for f in files:
66        basename_to_srcs_map[os.path.basename(f)].append(f)
67
68    duplicates_exist = False
69    for (basename, srcs) in basename_to_srcs_map.items():
70        if len(srcs) > 1:
71            duplicates_exist = True
72            logging.error('Destination filename "%s" has multiple possible sources: %s',
73                         basename, srcs)
74
75    if duplicates_exist:
76        sys.exit(1)
77
78
79def files_to_dist(pattern):
80    # Assume that dist.bzl is in the same package as dist.py
81    runfiles_directory = os.path.dirname(__file__)
82    dist_manifests = glob.glob(
83        os.path.join(runfiles_directory, pattern))
84    if not dist_manifests:
85        logging.warning("Could not find a file with pattern %s"
86                        " in the runfiles directory: %s", pattern, runfiles_directory)
87    files_to_dist = []
88    for dist_manifest in dist_manifests:
89        with open(dist_manifest, "r") as f:
90            files_to_dist += [line.strip() for line in f]
91    return files_to_dist
92
93
94def copy_files_to_dist_dir(files, archives, mode_overrides, dist_dir, flat, prefix,
95    strip_components, archive_prefix, wipe_dist_dir, allow_duplicate_filenames, **ignored):
96
97    if flat and not allow_duplicate_filenames:
98        ensure_unique_filenames(files)
99
100    if wipe_dist_dir and os.path.exists(dist_dir):
101        shutil.rmtree(dist_dir)
102
103    logging.info("Copying to %s", dist_dir)
104
105    for src in files:
106        if flat:
107            src_relpath = os.path.basename(src)
108        elif strip_components > 0:
109            src_relpath = src.split('/', strip_components)[-1]
110        else:
111            src_relpath = src
112
113        src_relpath = os.path.join(prefix, src_relpath)
114
115        dst = os.path.join(dist_dir, src_relpath)
116        if os.path.isfile(src):
117            dst_dirname = os.path.dirname(dst)
118            logging.debug("Copying file: %s" % dst)
119            if not os.path.exists(dst_dirname):
120                os.makedirs(dst_dirname)
121
122            copy_with_modes(src, dst, mode_overrides)
123        elif os.path.isdir(src):
124            logging.debug("Copying dir: %s" % dst)
125            if os.path.exists(dst):
126                # make the directory temporary writable, then
127                # shutil.copytree will restore correct permissions.
128                os.chmod(dst, 750)
129            shutil.copytree(
130                os.path.abspath(src),
131                dst,
132                copy_function=lambda s, d: copy_with_modes(s, d, mode_overrides),
133                dirs_exist_ok=True,
134            )
135
136    for archive in archives:
137        try:
138            with tarfile.open(archive) as tf:
139                dst_dirname = os.path.join(dist_dir, archive_prefix)
140                logging.debug("Extracting archive: %s -> %s", archive, dst_dirname)
141                tf.extractall(dst_dirname)
142        except tarfile.TarError:
143            # toybox does not support creating empty tar files, hence the build
144            # system may use empty files as empty archives.
145            if os.path.getsize(archive) == 0:
146                logging.warning("Skipping empty tar file: %s", archive)
147                continue
148             # re-raise if we do not know anything about this error
149            logging.exception("Unknown TarError.")
150            raise
151
152
153def config_logging(log_level_str):
154    level = getattr(logging, log_level_str.upper(), None)
155    if not isinstance(level, int):
156        sys.stderr.write("ERROR: Invalid --log {}\n".format(log_level_str))
157        sys.exit(1)
158    logging.basicConfig(level=level, format="[dist] %(levelname)s: %(message)s")
159
160
161def main():
162    parser = argparse.ArgumentParser(
163        description="Dist Bazel output files into a custom directory.")
164    parser.add_argument(
165        "--dist_dir", required=True, help="""path to the dist dir.
166            If relative, it is interpreted as relative to Bazel workspace root
167            set by the BUILD_WORKSPACE_DIRECTORY environment variable, or
168            PWD if BUILD_WORKSPACE_DIRECTORY is not set.""")
169    parser.add_argument(
170        "--flat",
171        action="store_true",
172        help="ignore subdirectories in the manifest")
173    parser.add_argument(
174        "--strip_components", type=int, default=0,
175        help="number of leading components to strip from paths before applying --prefix")
176    parser.add_argument(
177        "--prefix", default="",
178        help="path prefix to apply within dist_dir for copied files")
179    parser.add_argument(
180        "--archive_prefix", default="",
181        help="Path prefix to apply within dist_dir for extracted archives. " +
182             "Supported archives: tar.")
183    parser.add_argument("--log", help="Log level (debug, info, warning, error)", default="debug")
184    parser.add_argument(
185        "--wipe_dist_dir",
186        action="store_true",
187        help="remove existing dist_dir prior to running"
188    )
189    parser.add_argument(
190        "--allow_duplicate_filenames",
191        action="store_true",
192        help="allow multiple files with the same name to be copied to dist_dir (overwriting)"
193    )
194    parser.add_argument(
195        "--mode_override",
196        metavar=("PATTERN", "MODE"),
197        action="append",
198        nargs=2,
199        default=[],
200        help='glob pattern and mode to set on files matching pattern (e.g. --mode_override "*.sh" "755")'
201    )
202
203    args = parser.parse_args(sys.argv[1:])
204
205    mode_overrides = []
206    for (pattern, mode) in args.mode_override:
207        try:
208            mode_overrides.append((pattern, int(mode, 8)))
209        except ValueError:
210            logging.error("invalid octal permissions: %s", mode)
211            sys.exit(1)
212
213    config_logging(args.log)
214
215    if not os.path.isabs(args.dist_dir):
216        # BUILD_WORKSPACE_DIRECTORY is the root of the Bazel workspace containing
217        # this binary target.
218        # https://docs.bazel.build/versions/main/user-manual.html#run
219        args.dist_dir = os.path.join(
220            os.environ.get("BUILD_WORKSPACE_DIRECTORY"), args.dist_dir)
221
222    files = files_to_dist("*_dist_manifest.txt")
223    archives = files_to_dist("*_dist_archives_manifest.txt")
224    copy_files_to_dist_dir(files, archives, mode_overrides, **vars(args))
225
226
227if __name__ == "__main__":
228    main()
229