utils/lost_my_space.py at master · ttsiodras/utils · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
#!/usr/bin/env python3
"""
Keep track of used storage in my filesystem.
Launch this once, with "-u" - and it will create a zstd-compressed
snapshot of the files in your current filesystem (and their sizes).

Launch it again, and it will show you the difference between the
previous snapshot and your current state.
"""
import os
import sys
import pickle
import subprocess
from stat import S_ISDIR


def visit_fs(action):
    """
    Used for both storing a snapshot and comparing against one.
    Calls action with each and every filename in your filesystem.
    """
    try:
        exclusions = {
            x.strip()
            for x in open(
                os.path.dirname(os.path.abspath(__file__)) +
                os.sep +
                "lost_my_space.exceptions")
        }
    except:
        exclusions = {}
    for pth, dirlist, filelist in os.walk("/"):
        if any(pth.startswith(x) for x in exclusions):
            continue
        dirlist[:] = [
            x for x in dirlist
            if not os.path.ismount(pth + os.sep + x)]
        for file in filelist:
            action(pth + os.sep + file)


def main():
    """
    Store - or compare against - a filesystem snapshot.
    """
    if "-u" in sys.argv:
        all_files, cnt = {}, [0]
        print("[-] Files indexed:            ", flush=True, end='',
              file=sys.stderr)

        def action_store(fullpath):
            try:
                statdata = os.lstat(fullpath)
                if not S_ISDIR(statdata.st_mode):
                    size = statdata.st_size
                    all_files[fullpath] = size
                    cnt[0] += 1
                    if cnt[0] & 0xFFFF == 0:
                        print("\b\b\b\b\b\b\b\b\b\b\b%11d" % cnt[0],
                              flush=True, end='', file=sys.stderr)
            except Exception:
                pass
        visit_fs(action_store)
        print("\b\b\b\b\b\b\b\b\b\b\b%11d" % cnt[0], flush=True,
              file=sys.stderr)
        pickle.dump(all_files, open("all_files.pickle", "wb"))
        print("[-] Compressing DB via zstd...", file=sys.stderr)
        os.system("zstd --rm -f -9 all_files.pickle")
    else:
        print("[-] Reading previous filesystem snapshot...",
              file=sys.stderr)
        cmd = ['zstdcat', 'all_files.pickle.zst']
        proc = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, _ = proc.communicate()
        all_files = pickle.loads(stdout)
        changes = [0, 0]

        def action_compare(fullpath):
            try:
                statdata = os.lstat(fullpath)
                if not S_ISDIR(statdata.st_mode):
                    size = statdata.st_size
                    try:
                        if size == all_files[fullpath]:
                            return
                        delta = size-all_files[fullpath]
                        print("%11d [C] (%d) %s" % (size, delta, fullpath))
                        changes[0] += delta
                    except KeyError:
                        print("%11d [N] %s" % (size, fullpath))
                        changes[1] += size
            except Exception:
                pass
        print("[-] Comparing with current...", file=sys.stderr)
        visit_fs(action_compare)
        print("[-] Due to new files:      %d" % changes[1], file=sys.stderr)
        print("[-] Due to modified files: %d" % changes[0], file=sys.stderr)

if __name__ == "__main__":
    main()