1#!/usr/bin/python3
2
3import argparse
4import os
5import shutil
6import subprocess
7import sys
8import tempfile
9from multiprocessing import Pool, cpu_count
10from os.path import getsize
11
12# Set command line arguments
13parser = argparse.ArgumentParser(
14    formatter_class=argparse.ArgumentDefaultsHelpFormatter
15)
16
17parser.add_argument(
18    "-b",
19    "--build_dir",
20    dest="BUILD_DIR",
21    default="/home/ed/openbmc-openbmc",
22    help="Build directory path.",
23)
24
25parser.add_argument(
26    "-s",
27    "--squashfs_file",
28    dest="SQUASHFS_FILE",
29    default="/build/tmp/deploy/images/wolfpass"
30    + "/intel-platforms-wolfpass.squashfs-xz",
31    help="Squashfs file.",
32)
33
34parser.add_argument(
35    "-t",
36    "--threads",
37    dest="threads",
38    default=int(cpu_count()),
39    type=int,
40    help="Number of threads to use (defaults to cpu count)",
41)
42
43args = parser.parse_args()
44
45# files below this size wont be attempted
46FILE_SIZE_LIMIT = 0
47
48SQUASHFS = args.SQUASHFS_FILE
49if not os.path.isabs(args.SQUASHFS_FILE):
50    SQUASHFS = args.BUILD_DIR + args.SQUASHFS_FILE
51
52original_size = getsize(SQUASHFS)
53print("squashfs size: {}".format(original_size))
54
55results = []
56
57
58def get_unsquash_results(filepath):
59    with tempfile.TemporaryDirectory() as newsquashfsroot:
60        input_path = os.path.join(newsquashfsroot, "input")
61        shutil.copytree(
62            squashfsdir,
63            input_path,
64            symlinks=True,
65            ignore_dangling_symlinks=True,
66        )
67        file_to_remove = os.path.join(input_path, filepath)
68        try:
69            os.remove(file_to_remove)
70        except IsADirectoryError:
71            shutil.rmtree(file_to_remove)
72        subprocess.check_output(
73            [
74                "mksquashfs",
75                input_path,
76                newsquashfsroot + "/test",
77                "-comp",
78                "xz",
79                "-processors",
80                "1",
81            ]
82        )
83
84        return (
85            filepath.replace(squashfsdir, ""),
86            original_size - getsize(newsquashfsroot + "/test"),
87        )
88
89
90with tempfile.TemporaryDirectory() as tempsquashfsdir:
91    print("writing to " + tempsquashfsdir)
92    squashfsdir = os.path.join(tempsquashfsdir, "squashfs-root")
93    # squashfsdir = os.path.join("/tmp", "squashfs-root")
94    command = ["unsquashfs", "-d", squashfsdir, SQUASHFS]
95    print(" ".join(command))
96    subprocess.check_call(command)
97
98    files_to_test = []
99    for root, dirs, files in os.walk(squashfsdir):
100        for name in files + dirs:
101            filepath = os.path.join(root, name)
102            if not os.path.islink(filepath):
103                if getsize(filepath) > FILE_SIZE_LIMIT:
104                    files_to_test.append(
105                        os.path.relpath(filepath, squashfsdir)
106                    )
107
108    print("{} files to attempt removing".format(len(files_to_test)))
109
110    print("Using {} threads".format(args.threads))
111    with Pool(args.threads) as p:
112        for i, res in enumerate(
113            p.imap_unordered(get_unsquash_results, files_to_test)
114        ):
115            results.append(res)
116            sys.stderr.write(
117                "\rdone {:.1f}%".format(100 * (i / len(files_to_test)))
118            )
119
120results.sort(key=lambda x: x[1], reverse=True)
121
122with open("results.txt", "w") as result_file:
123    for filepath, size in results:
124        result = "{:>10}: {}".format(size, filepath)
125        print(result)
126        result_file.write(result + "\n")
127