Completed
Push — master ( 654bda...c0db53 )
by John
03:04
created

dirsizer()   A

Complexity

Conditions 2

Size

Total Lines 17

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 2
CRAP Score 2

Importance

Changes 0
Metric Value
cc 2
c 0
b 0
f 0
dl 0
loc 17
ccs 2
cts 2
cp 1
crap 2
rs 9.4285
1
#!/usr/bin/env python3
2 4
"""This module is used to operate with bar files."""
3
4 4
import os  # filesystem read
5 4
import zipfile  # zip extract, zip compresssion
6 4
import shutil  # folder operations
7 4
import base64  # encoding for hashes
8 4
import hashlib   # get hashes
9 4
from bbarchivist import utilities  # platform determination
10 4
from bbarchivist import exceptions  # exception handling
11 4
from bbarchivist import bbconstants  # premade stuff
12
13 4
__author__ = "Thurask"
14 4
__license__ = "WTFPL v2"
15 4
__copyright__ = "Copyright 2015-2016 Thurask"
16
17
18 4
def extract_bars(filepath):
19
    """
20
    Extract .signed files from .bar files.
21
    Use system zlib.
22
23
    :param filepath: Path to bar file directory.
24
    :type filepath: str
25
    """
26 4
    try:
27 4
        for file in os.listdir(filepath):
28 4
            extract_individual_bar(file, filepath)
29 4
    except (RuntimeError, OSError) as exc:
30 4
        exceptions.handle_exception(exc, "EXTRACTION FAILURE", None)
31
32
33 4
def extract_individual_bar(file, filepath):
34
    """
35
    Generate bar file contents and extract signed files.
36
37
    :param file: Bar file to extract.
38
    :type file: str
39
40
    :param filepath: Path to bar file directory.
41
    :type filepath: str
42
    """
43 4
    if file.endswith(".bar"):
44 4
        print("EXTRACTING: {0}".format(file))
45 4
        zfile = zipfile.ZipFile(os.path.join(filepath, file), 'r')
46 4
        names = zfile.namelist()
47 4
        extract_signed_file(zfile, names, filepath)
48
49
50 4
def extract_signed_file(zfile, names, filepath):
51
    """
52
    Extract signed file from a provided bar.
53
54
    :param zfile: Open (!!!) ZipFile instance.
55
    :type zfile: zipfile.ZipFile
56
57
    :param names: List of bar file contents.
58
    :type names: list(str)
59
60
    :param filepath: Path to bar file directory.
61
    :type filepath: str
62
    """
63 4
    for name in names:
64 4
        if str(name).endswith(".signed"):
65 4
            zfile.extract(name, filepath)
66
67
68 4
def get_sha512_manifest(zfile):
69
    """
70
    Get MANIFEST.MF from a bar file.
71
72
    :param zfile: Open (!!!) ZipFile instance.
73
    :type zfile: zipfile.ZipFile
74
    """
75 4
    names = zfile.namelist()
76 4
    manifest = None
77 4
    for name in names:
78 4
        if name.endswith("MANIFEST.MF"):
79 4
            manifest = name
80 4
            break
81 4
    if manifest is None:
82 4
        raise SystemExit
83 4
    return manifest
84
85
86 4
def get_sha512_from_manifest(manf):
87
    """
88
    Retrieve asset name and hash from MANIFEST.MF file.
89
90
    :param manf: Content of MANIFEST.MF file, in bytes.
91
    :type manf: list(bytes)
92
    """
93 4
    alist = []
94 4
    for idx, line in enumerate(manf):
95 4
        if line.endswith(b"signed"):
96 4
            alist.append(manf[idx])
97 4
            alist.append(manf[idx + 1])
98 4
    assetname = alist[0].split(b": ")[1]
99 4
    assethash = alist[1].split(b": ")[1]
100 4
    return assetname, assethash
101
102
103 4
def retrieve_sha512(filename):
104
    """
105
    Get the premade, Base64 encoded SHA512 hash of a signed file in a bar.
106
107
    :param filename: Bar file to check.
108
    :type filename: str
109
    """
110 4
    try:
111 4
        zfile = zipfile.ZipFile(filename, 'r')
112 4
        manifest = get_sha512_manifest(zfile)
113 4
        manf = zfile.read(manifest).splitlines()
114 4
        assetname, assethash = get_sha512_from_manifest(manf)
115 4
        return assetname, assethash  # (b"blabla.signed", b"somehash")
116 4
    except (RuntimeError, OSError, zipfile.BadZipFile) as exc:
117 4
        exceptions.handle_exception(exc, "EXTRACTION FAILURE", None)
118
119
120 4
def verify_sha512(filename, inithash):
121
    """
122
    Compare the original hash value with the current.
123
124
    :param filename: Signed file to check.
125
    :type filename: str
126
127
    :param inithash: Original SHA512 hash, as bytestring.
128
    :type inithash: bytes
129
    """
130 4
    sha512 = hashlib.sha512()
131 4
    with open(filename, 'rb') as file:
132 4
        while True:
133 4
            data = file.read(16 * 1024 * 1024)
134 4
            if not data:
135 4
                break
136 4
            sha512.update(data)
137 4
    rawdigest = sha512.digest()  # must be bytestring, not hexadecimalized str
138 4
    b64h = base64.b64encode(rawdigest, altchars=b"-_")  # replace some chars
139 4
    b64h = b64h.strip(b"==")  # remove padding
140 4
    return b64h == inithash
141
142
143 4
def bar_tester(filepath):
144
    """
145
    Use zipfile in order to test a bar for errors.
146
147
    :param filepath: Path to bar file.
148
    :type filepath: str
149
    """
150 4
    try:
151 4
        with zipfile.ZipFile(filepath, "r") as zfile:
152 4
            brokens = zfile.testzip()
153 4
    except zipfile.BadZipFile:
154 4
        brokens = filepath
155 4
    return brokens
156
157
158 4
def remove_empty_folders(a_folder):
159
    """
160
    Remove empty folders in a given folder using os.walk().
161
162
    :param a_folder: Target folder.
163
    :type a_folder: str
164
    """
165 4
    for curdir, subdirs, files in os.walk(a_folder):
166 4
        while True:
167 4
            try:
168 4
                if not subdirs and not files:
169 4
                    os.rmdir(curdir)
170 4
            except OSError:
171
                continue
172 4
            except NotImplementedError:
173 4
                break
174 4
            break
175
176
177 4
def persistent_remove(afile):
178
    """
179
    Remove a file, and if it doesn't want to remove, keep at it.
180
181
    :param afile: Path to file you want terminated with extreme prejudice.
182
    :type afile: str
183
    """
184 4
    while True:
185 4
        try:
186 4
            os.remove(afile)
187
        except PermissionError:
188
            continue
189
        else:
190 4
            break
191
192
193 4
def remove_signed_files(a_folder):
194
    """
195
    Remove signed files from a given folder.
196
197
    :param a_folder: Target folder.
198
    :type a_folder: str
199
    """
200 4
    files = [os.path.abspath(os.path.join(a_folder, file)) for file in os.listdir(a_folder)]
201 4
    for afile in files:
202 4
        if afile.endswith(".signed") and os.path.exists(afile):
203 4
            print("REMOVING: {0}".format(os.path.basename(afile)))
204 4
            persistent_remove(afile)
205
206
207 4
def remove_unpacked_loaders(osdir, raddir, radios):
208
    """
209
    Remove uncompressed loader folders.
210
211
    :param osdir: OS loader folder.
212
    :type osdir: str
213
214
    :param raddir: Radio loader folder.
215
    :type raddir: str
216
217
    :param radios: If we made radios this run.
218
    :type radios: bool
219
    """
220 4
    utilities.cond_do(shutil.rmtree, [osdir, raddir], condition=radios)
221
222
223 4
def create_blitz(a_folder, swver):
224
    """
225
    Create a blitz file: a zipped archive of all app/core/radio bars.
226
227
    :param a_folder: Target folder.
228
    :type a_folder: str
229
230
    :param swver: Software version to title the blitz.
231
    :type swver: str
232
    """
233 4
    fname = "Blitz-{0}.zip".format(swver)
234 4
    with zipfile.ZipFile(fname, 'w', zipfile.ZIP_DEFLATED, allowZip64=True) as zfile:
235 4
        for root, dirs, files in os.walk(a_folder):
236 4
            del dirs
237 4
            for file in files:
238 4
                print("ZIPPING: {0}".format(utilities.stripper(file)))
239 4
                abs_filename = os.path.join(root, file)
240 4
                abs_arcname = os.path.basename(abs_filename)
241 4
                zfile.write(abs_filename, abs_arcname)
242
243
244 4
def move_loaders(ldir,
245
                 exedir_os, exedir_rad,
246
                 zipdir_os, zipdir_rad):
247
    """
248
    Move autoloaders to zipped and loaders directories in localdir.
249
250
    :param ldir: Local directory, containing files you wish to move.
251
    :type ldir: str
252
253
    :param exedir_os: Large autoloader .exe destination.
254
    :type exedir_os: str
255
256
    :param exedir_rad: Small autoloader .exe destination.
257
    :type exedir_rad: str
258
259
    :param zipdir_os: Large autoloader archive destination.
260
    :type zipdir_os: str
261
262
    :param zipdir_rad: Small autoloader archive destination.
263
    :type zipdir_rad: str
264
    """
265 4
    arx = bbconstants.ARCS
266 4
    pfx = bbconstants.PREFIXES
267 4
    loaders = [os.path.join(ldir, file) for file in os.listdir(ldir) if utilities.prepends(file, pfx, ".exe")]
268 4
    move_loader_pairs(loaders, exedir_os, exedir_rad)
269 4
    zippeds = [os.path.join(ldir, file) for file in os.listdir(ldir) if utilities.prepends(file, pfx, arx)]
270 4
    move_loader_pairs(zippeds, zipdir_os, zipdir_rad)
271
272
273 4
def move_loader_pairs(files, dir_os, dir_rad):
274
    """
275
    Move autoloaders to zipped/loaders directories.
276
277
    :param files: List of autoloader files.
278
    :type files: list(str)
279
280
    :param dir_os: Large autoloader destination.
281
    :type dir_os: str
282
283
    :param dir_rad: Small autoloader destination.
284
    :type dir_rad: str
285
    """
286 4
    for file in files:
287 4
        print("MOVING: {0}".format(os.path.basename(file)))
288 4
        dest_os = os.path.join(dir_os, os.path.basename(file))
289 4
        dest_rad = os.path.join(dir_rad, os.path.basename(file))
290 4
        loader_sorter(file, dest_os, dest_rad)
291
292
293 4
def dirsizer(file, osdir, raddir, maxsize=90 * 1000 * 1000):
294
    """
295
    Return output directory based in input filesize.
296
297
    :param file: The file to sort. Absolute paths, please.
298
    :type file: str
299
300
    :param osdir: Large file destination.
301
    :type osdir: str
302
303
    :param raddir: Small file destination.
304
    :type raddir: str
305
306
    :param maxsize: Return osdir if filesize > maxsize else raddir. Default is 90MB.
307
    :type maxsize: int
308
    """
309 4
    return osdir if os.path.getsize(file) > maxsize else raddir
310
311
312 4
def loader_sorter(file, osdir, raddir):
313
    """
314
    Sort loaders based on size.
315
316
    :param file: The file to sort. Absolute paths, please.
317
    :type file: str
318
319
    :param osdir: Large file destination.
320
    :type osdir: str
321
322
    :param raddir: Small file destination.
323
    :type raddir: str
324
    """
325 4
    outdir = dirsizer(file, osdir, raddir)
326 4
    persistent_move(file, outdir)
327
328
329 4
def move_bars(localdir, osdir, radiodir):
330
    """
331
    Move bar files to subfolders of a given folder.
332
333
    :param localdir: Directory to use.
334
    :type localdir: str
335
336
    :param osdir: OS file directory (large bars).
337
    :type osdir: str
338
339
    :param radiodir: Radio file directory (small bars).
340
    :type radiodir: str
341
    """
342 4
    for files in os.listdir(localdir):
343 4
        if files.endswith(".bar"):
344 4
            print("MOVING: {0}".format(files))
345 4
            herefile = os.path.join(localdir, files)
0 ignored issues
show
Unused Code introduced by
The variable herefile seems to be unused.
Loading history...
346 4
            outdir = dirsizer(files, osdir, radiodir)
347 4
            atomic_move(files, outdir)
348
349
350 4
def persistent_move(infile, outdir):
351
    """
352
    Move file to given folder, removing file if it exists in folder.
353
354
    :param infile: Path to file to move.
355
    :type infile: str
356
357
    :param outdir: Directory to move to.
358
    :type outdir: str
359
    """
360 4
    while True:
361 4
        try:
362 4
            shutil.move(infile, outdir)
363
        except shutil.Error:
364
            os.remove(infile)
365
            continue
366 4
        break
367
368
369 4
def atomic_move(infile, outdir):
370
    """
371
    Move file to given folder, removing if things break.
372
373
    :param infile: Path to file to move.
374
    :type infile: str
375
376
    :param outdir: Directory to move to.
377
    :type outdir: str
378
    """
379 4
    try:
380 4
        shutil.move(infile, outdir)
381
    except shutil.Error:
382
        os.remove(os.path.join(outdir, infile))
383
384
385 4
def replace_bar_pair(localdir, osfile, radfile):
386
    """
387
    Move pair of OS and radio bars to a given folder.
388
389
    :param localdir: Final bar directory.
390
    :type localdir: str
391
392
    :param osfile: Path to OS file.
393
    :type osfile: str
394
395
    :param radfile: Path to radio file.
396
    :type radfile: str
397
    """
398 4
    shutil.move(osfile, localdir)
399 4
    shutil.move(radfile, localdir)
400
401
402 4
def replace_bars_bulk(localdir, barfiles):
403
    """
404
    Move set of OS and radio bars to a given folder.
405
406
    :param localdir: Final bar directory.
407
    :type localdir: str
408
409
    :param barfiles: List of OS/radio file paths.
410
    :type barfiles: list(str)
411
    """
412 4
    for barfile in barfiles:
413 4
        shutil.move(barfile, os.path.abspath(localdir))
414
415
416 4
def make_folder(localdir, root):
417
    """
418
    Make a folder if it doesn't exist.
419
420
    :param localdir: Top level folder.
421
    :type localdir: str
422
423
    :param root: Folder to create.
424
    :type root: str
425
    """
426 4
    if not os.path.exists(os.path.join(localdir, root)):
427 4
        os.makedirs(os.path.join(localdir, root), exist_ok=True)
428 4
    return os.path.join(localdir, root)
429
430
431 4
def make_dirpairs(localdir, root, osversion, radioversion):
432
    """
433
    Create a pair of directories, with OS/radio versions for names.
434
435
    :param localdir: Top level folder.
436
    :type localdir: str
437
438
    :param root: Name for folder containing OS/radio pairs.
439
    :type root: str
440
441
    :param osversion: OS version.
442
    :type osversion: str
443
444
    :param radioversion: Radio version.
445
    :type radioversion: str
446
    """
447 4
    rootdir = make_folder(localdir, root)
448 4
    osdir = make_folder(rootdir, osversion)
449 4
    radiodir = make_folder(rootdir, radioversion)
450 4
    return osdir, radiodir
451
452
453 4
def make_dirs(localdir, osversion, radioversion):
454
    """
455
    Create the directory tree needed for archivist/lazyloader.
456
457
    :param localdir: Root folder.
458
    :type localdir: str
459
460
    :param osversion: OS version.
461
    :type osversion: str
462
463
    :param radioversion: Radio version.
464
    :type radioversion: str
465
    """
466 4
    os.makedirs(localdir, exist_ok=True)
467 4
    bardir_os, bardir_radio = make_dirpairs(localdir, "bars", osversion, radioversion)
468 4
    loaderdir_os, loaderdir_radio = make_dirpairs(localdir, "loaders", osversion, radioversion)
469 4
    zipdir_os, zipdir_radio = make_dirpairs(localdir, "zipped", osversion, radioversion)
470
    return (bardir_os, bardir_radio, loaderdir_os, loaderdir_radio, zipdir_os, zipdir_radio)
471