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