Completed
Push — master ( c0d006...1a22fd )
by John
01:20
created

smart_is_tarfile()   A

Complexity

Conditions 3

Size

Total Lines 13

Duplication

Lines 0
Ratio 0 %
Metric Value
cc 3
dl 0
loc 13
rs 9.4285
1
#!/usr/bin/env python3
2
"""This module is used to operate with bar files and other archives."""
3
4
import os  # filesystem read
5
import subprocess  # invocation of 7z, cap
6
import zipfile  # zip extract, zip compresssion
7
import tarfile  # txz/tbz/tgz compression
8
import sys  # version info
9
import shutil  # folder operations
10
import base64  # encoding for hashes
11
import hashlib   # get hashes
12
import configparser  # config parsing, duh
13
from bbarchivist import utilities  # platform determination
14
from bbarchivist import bbconstants  # premade stuff
15
from bbarchivist import decorators  # timer
16
17
__author__ = "Thurask"
18
__license__ = "WTFPL v2"
19
__copyright__ = "Copyright 2015-2016 Thurask"
20
21
22
def extract_bars(filepath):
23
    """
24
    Extract .signed files from .bar files.
25
    Use system zlib.
26
27
    :param filepath: Path to bar file directory.
28
    :type filepath: str
29
    """
30
    try:
31
        for file in os.listdir(filepath):
32
            if file.endswith(".bar"):
33
                print("EXTRACTING: {0}".format(file))
34
                zfile = zipfile.ZipFile(file, 'r')
35
                names = zfile.namelist()
36
                for name in names:
37
                    if str(name).endswith(".signed"):
38
                        zfile.extract(name, filepath)
39
    except (RuntimeError, OSError) as exc:
40
        print("EXTRACTION FAILURE")
41
        print(str(exc))
42
        print("DID IT DOWNLOAD PROPERLY?")
43
44
45
def retrieve_sha512(filename):
46
    """
47
    Get the premade, Base64 encoded SHA512 hash of a signed file in a bar.
48
49
    :param filename: Bar file to check.
50
    :type filename: str
51
    """
52
    try:
53
        zfile = zipfile.ZipFile(filename, 'r')
54
        names = zfile.namelist()
55
        manifest = None
56
        for name in names:
57
            if name.endswith("MANIFEST.MF"):
58
                manifest = name
59
                break
60
        if manifest is None:
61
            raise SystemExit
62
        manf = zfile.read(manifest).splitlines()
63
        alist = []
64
        for idx, line in enumerate(manf):
65
            if line.endswith(b"signed"):
66
                alist.append(manf[idx])
67
                alist.append(manf[idx + 1])
68
        assetname = alist[0].split(b": ")[1]
69
        assethash = alist[1].split(b": ")[1]
70
        return assetname, assethash  # (b"blabla.signed", b"somehash")
71
    except (RuntimeError, OSError, zipfile.BadZipFile) as exc:
72
        print("EXTRACTION FAILURE")
73
        print(str(exc))
74
        print("DID IT DOWNLOAD PROPERLY?")
75
76
77
def verify_sha512(filename, inithash):
78
    """
79
    Compare the original hash value with the current.
80
81
    :param filename: Signed file to check.
82
    :type filename: str
83
84
    :param inithash: Original SHA512 hash, as bytestring.
85
    :type inithash: bytes
86
    """
87
    sha512 = hashlib.sha512()
88
    with open(filename, 'rb') as file:
89
        while True:
90
            data = file.read(16 * 1024 * 1024)
91
            if not data:
92
                break
93
            sha512.update(data)
94
    rawdigest = sha512.digest()  # must be bytestring, not hexadecimalized str
95
    b64h = base64.b64encode(rawdigest, altchars=b"-_")  # replace some chars
96
    b64h = b64h.strip(b"==")  # remove padding
97
    return b64h == inithash
98
99
100
def bar_tester(filepath):
101
    """
102
    Use zipfile in order to test a bar for errors.
103
104
    :param filepath: Path to bar file.
105
    :type filepath: str
106
    """
107
    try:
108
        with zipfile.ZipFile(filepath, "r") as zfile:
109
            brokens = zfile.testzip()
110
    except zipfile.BadZipFile:
111
        brokens = filepath
112
    return brokens
113
114
115
def smart_is_tarfile(filepath):
116
    """
117
    :func:`tarfile.is_tarfile` plus error handling.
118
119
    :param filepath: Filename.
120
    :type filepath: str
121
    """
122
    try:
123
        istar = tarfile.is_tarfile(filepath)
124
    except (FileNotFoundError, OSError):
125
        return False
126
    else:
127
        return istar
128
129
130
@decorators.timer
131
def sz_compress(filepath, filename, szexe=None, strength=5, errors=False):
132
    """
133
    Pack a file into a LZMA2 7z file.
134
135
    :param filepath: Basename of file, no extension.
136
    :type filepath: str
137
138
    :param filename: Name of file to pack.
139
    :type filename: str
140
141
    :param szexe: Path to 7z executable.
142
    :type szexe: str
143
144
    :param strength: Compression strength. 5 is normal, 9 is ultra.
145
    :type strength: int
146
147
    :param errors: Print completion status message. Default is false.
148
    :type errors: bool
149
    """
150
    szcodes = {
151
        0: "NO ERRORS",
152
        1: "COMPLETED WITH WARNINGS",
153
        2: "FATAL ERROR",
154
        7: "COMMAND LINE ERROR",
155
        8: "OUT OF MEMORY ERROR",
156
        255: "USER STOPPED PROCESS"
157
    }
158
    strength = str(strength)
159
    rawname = os.path.dirname(filepath)
160
    thr = str(utilities.get_core_count())
161
    fold = os.path.join(rawname, filename)
162
    cmd = '{0} a -mx{1} -m0=lzma2 -mmt{2} "{3}.7z" "{4}"'.format(szexe, strength,
163
                                                                 thr, filepath, fold)
164
    with open(os.devnull, 'wb') as dnull:
165
        excode = subprocess.call(cmd, stdout=dnull, stderr=subprocess.STDOUT, shell=True)
166
    if errors:
167
        print(szcodes[excode])
168
169
170
def sz_verify(filepath, szexe=None):
171
    """
172
    Verify that a .7z file is valid and working.
173
174
    :param filepath: Filename.
175
    :type filepath: str
176
177
    :param szexe: Path to 7z executable.
178
    :type szexe: str
179
    """
180
    filepath = os.path.abspath(filepath)
181
    cmd = '{0} t "{1}"'.format(szexe, filepath)
182
    with open(os.devnull, 'wb') as dnull:
183
        excode = subprocess.call(cmd, stdout=dnull, stderr=subprocess.STDOUT, shell=True)
184
    return excode == 0
185
186
187
@decorators.timer
188
def tar_compress(filepath, filename):
189
    """
190
    Pack a file into an uncompressed tarfile.
191
192
    :param filepath: Basename of file, no extension.
193
    :type filepath: str
194
195
    :param filename: Name of file to pack.
196
    :type filename: str
197
    """
198
    with tarfile.open("{0}.tar".format(filepath), 'w:') as tfile:
199
        tfile.add(filename, filter=None)
200
201
202
def tar_verify(filepath):
203
    """
204
    Verify that a tar file is valid and working.
205
206
    :param filepath: Filename.
207
    :type filepath: str
208
    """
209
    if smart_is_tarfile(filepath):
210
        with tarfile.open(filepath, "r:") as thefile:
211
            mems = thefile.getmembers()
212
        return False if not mems else True
213
    else:
214
        return False
215
216
217
@decorators.timer
218
def tgz_compress(filepath, filename, strength=5):
219
    """
220
    Pack a file into a gzip tarfile.
221
222
    :param filepath: Basename of file, no extension.
223
    :type filepath: str
224
225
    :param filename: Name of file to pack.
226
    :type filename: str
227
228
    :param strength: Compression strength. 5 is normal, 9 is ultra.
229
    :type strength: int
230
    """
231
    with tarfile.open("{0}.tar.gz".format(filepath), 'w:gz', compresslevel=strength) as gzfile:
232
        gzfile.add(filename, filter=None)
233
234
235
def tgz_verify(filepath):
236
    """
237
    Verify that a tar.gz file is valid and working.
238
239
    :param filepath: Filename.
240
    :type filepath: str
241
    """
242
    if smart_is_tarfile(filepath):
243
        with tarfile.open(filepath, "r:gz") as thefile:
244
            mems = thefile.getmembers()
245
        return False if not mems else True
246
    else:
247
        return False
248
249
250
@decorators.timer
251
def tbz_compress(filepath, filename, strength=5):
252
    """
253
    Pack a file into a bzip2 tarfile.
254
255
    :param filepath: Basename of file, no extension.
256
    :type filepath: str
257
258
    :param filename: Name of file to pack.
259
    :type filename: str
260
261
    :param strength: Compression strength. 5 is normal, 9 is ultra.
262
    :type strength: int
263
    """
264
    with tarfile.open("{0}.tar.bz2".format(filepath), 'w:bz2', compresslevel=strength) as bzfile:
265
        bzfile.add(filename, filter=None)
266
267
268
def tbz_verify(filepath):
269
    """
270
    Verify that a tar.bz2 file is valid and working.
271
272
    :param filepath: Filename.
273
    :type filepath: str
274
    """
275
    if smart_is_tarfile(filepath):
276
        with tarfile.open(filepath, "r:bz2") as thefile:
277
            mems = thefile.getmembers()
278
        return False if not mems else True
279
    else:
280
        return False
281
282
283
@decorators.timer
284
def txz_compress(filepath, filename):
285
    """
286
    Pack a file into a LZMA tarfile.
287
288
    :param filepath: Basename of file, no extension.
289
    :type filepath: str
290
291
    :param filename: Name of file to pack.
292
    :type filename: str
293
    """
294
    with tarfile.open("{0}.tar.xz".format(filepath), 'w:xz') as xzfile:
295
        xzfile.add(filename, filter=None)
296
297
298
def txz_verify(filepath):
299
    """
300
    Verify that a tar.xz file is valid and working.
301
302
    :param filepath: Filename.
303
    :type filepath: str
304
    """
305
    if sys.version_info[1] < 3:
306
        pass
307
    else:
308
        if smart_is_tarfile(filepath):
309
            with tarfile.open(filepath, "r:xz") as thefile:
310
                mems = thefile.getmembers()
311
            return False if not mems else True
312
        else:
313
            return False
314
315
316
@decorators.timer
317
def zip_compress(filepath, filename):
318
    """
319
    Pack a file into a DEFLATE zipfile.
320
321
    :param filepath: Basename of file, no extension.
322
    :type filepath: str
323
324
    :param filename: Name of file to pack.
325
    :type filename: str
326
    """
327
    with zipfile.ZipFile("{0}.zip".format(filepath), 'w', zipfile.ZIP_DEFLATED, allowZip64=True) as zfile:
1 ignored issue
show
Coding Style introduced by
This line is too long as per the coding-style (106/100).

This check looks for lines that are too long. You can specify the maximum line length.

Loading history...
328
        zfile.write(filename)
329
330
331
def zip_verify(filepath):
332
    """
333
    Verify that a .zip file is valid and working.
334
335
    :param filepath: Filename.
336
    :type filepath: str
337
    """
338
    if zipfile.is_zipfile(filepath):
339
        brokens = bar_tester(filepath)
340
        return brokens != filepath
341
    else:
342
        return False
343
344
345
def filter_method(method, szexe=None):
346
    """
347
    Make sure methods are OK.
348
349
    :param method: Compression method to use.
350
    :type method: str
351
352
    :param szexe: Path to 7z executable, if needed.
353
    :type szexe: str
354
    """
355
    if sys.version_info[1] < 3 and method == "txz":
356
        method = "zip"  # fallback
357
    if method == "7z" and szexe is None:
358
        ifexists = utilities.prep_seven_zip()  # see if 7z exists
359
        if not ifexists:
360
            method = "zip"  # fallback
361
        else:
362
            szexe = utilities.get_seven_zip(False)
363
    return method
364
365
366
def calculate_strength():
367
    """
368
    Determine zip/gzip/bzip2 strength by OS bit setting.
369
    """
370
    strength = 9 if utilities.is_amd64() else 5
371
    return strength
372
373
374
def compress(filepath, method="7z", szexe=None, selective=False, errors=False):
375
    """
376
    Compress all autoloader files in a given folder, with a given method.
377
378
    :param filepath: Working directory. Required.
379
    :type filepath: str
380
381
    :param method: Compression type. Default is "7z".
382
    :type method: str
383
384
    :param szexe: Path to 7z executable, if needed.
385
    :type szexe: str
386
387
    :param selective: Only compress autoloaders. Default is false.
388
    :type selective: bool
389
390
    :param errors: Print completion status message. Default is false.
391
    :type errors: bool
392
    """
393
    method = filter_method(method, szexe)
394
    arx = bbconstants.ARCS
395
    pfx = bbconstants.PREFIXES
396
    files = [file for file in os.listdir(filepath) if not os.path.isdir(file)]
397
    if selective:
398
        filt0 = [file for file in files if prepends(file, pfx, "")]
399
        filt1 = [file for file in filt0 if not prepends(file, "", arx)]
400
        filt2 = [file for file in filt1 if prepends(file, "", ".exe")]
401
    else:
402
        filt2 = [file for file in files if not prepends(file, "", arx)]
403
    filt3 = [os.path.join(filepath, file) for file in filt2]
404
    for file in filt3:
405
        filename = os.path.splitext(os.path.basename(file))[0]
406
        fileloc = os.path.join(filepath, filename)
407
        print("COMPRESSING: {0}.exe".format(filename))
408
        if method == "7z":
409
            sz_compress(fileloc, file, szexe, calculate_strength(), errors)
410
        elif method == "tgz":
411
            tgz_compress(fileloc, file, calculate_strength())
412
        elif method == "txz":
413
            txz_compress(fileloc, file)
414
        elif method == "tbz":
415
            tbz_compress(fileloc, file, calculate_strength())
416
        elif method == "tar":
417
            tar_compress(fileloc, file)
418
        elif method == "zip":
419
            zip_compress(fileloc, file)
420
    return True
421
422
423
def verify(thepath, method="7z", szexe=None, selective=False):
424
    """
425
    Verify specific archive files in a given folder.
426
427
    :param thepath: Working directory. Required.
428
    :type thepath: str
429
430
    :param method: Compression type. Default is "7z". Defined in source.
431
    :type method: str
432
433
    :param szexe: Path to 7z executable, if needed.
434
    :type szexe: str
435
436
    :param selective: Only verify autoloaders. Default is false.
437
    :type selective: bool
438
    """
439
    method = filter_method(method, szexe)
440
    pathlist = [os.path.join(thepath, file) for file in os.listdir(thepath)]
441
    files = [file for file in pathlist if not os.path.isdir(file)]
442
    for file in files:
443
        filt = file.endswith(bbconstants.ARCS)
444
        if selective:
445
            filt = filt and file.startswith(bbconstants.PREFIXES)
446
        if filt:
447
            print("VERIFYING: {0}".format(file))
448
            if file.endswith(".7z") and szexe is not None:
449
                verif = sz_verify(os.path.abspath(file), szexe)
450
            elif file.endswith(".tar.gz"):
451
                verif = tgz_verify(file)
452
            elif file.endswith(".tar.xz"):
453
                verif = txz_verify(file)
454
            elif file.endswith(".tar.bz2"):
455
                verif = tbz_verify(file)
456
            elif file.endswith(".tar"):
457
                verif = tar_verify(file)
458
            elif file.endswith(".zip"):
459
                verif = zip_verify(file)
460
            if not verif:
461
                print("{0} IS BROKEN!".format((file)))
462
            return verif
463
464
465
def compress_suite(filepath, method="7z", szexe=None, selective=False):
466
    """
467
    Wrap compression and verification into one.
468
469
    :param filepath: Working directory. Required.
470
    :type filepath: str
471
472
    :param method: Compression type. Default is "7z". Defined in source.
473
    :type method: str
474
475
    :param szexe: Path to 7z executable, if needed.
476
    :type szexe: str
477
478
    :param selective: Only compress autoloaders. Default is false.
479
    :type selective: bool
480
    """
481
    compress(filepath, method, szexe, selective)
482
    verify(filepath, method, szexe, selective)
483
484
485
def remove_empty_folders(a_folder):
486
    """
487
    Remove empty folders in a given folder using os.walk().
488
489
    :param a_folder: Target folder.
490
    :type a_folder: str
491
    """
492
    for curdir, subdirs, files in os.walk(a_folder):
493
        while True:
494
            try:
495
                if not subdirs and not files:
496
                    os.rmdir(curdir)
497
            except OSError:
498
                continue
499
            except NotImplementedError:
500
                break
501
            break
502
503
504
def remove_signed_files(a_folder):
505
    """
506
    Remove signed files from a given folder.
507
508
    :param a_folder: Target folder.
509
    :type a_folder: str
510
    """
511
    files = [os.path.join(a_folder, file) for file in os.listdir(a_folder)]
512
    files = [os.path.abspath(file) for file in files]
513
    for file in files:
514
        if file.endswith(".signed") and os.path.exists(file):
515
            print("REMOVING: {0}".format(os.path.basename(file)))
516
            while True:
517
                try:
518
                    os.remove(os.path.abspath(file))
519
                except PermissionError:
520
                    continue
521
                else:
522
                    break
523
524
525
def remove_unpacked_loaders(osdir, raddir, radios):
526
    """
527
    Remove uncompressed loader folders.
528
529
    :param osdir: OS loader folder.
530
    :type osdir: str
531
532
    :param raddir: Radio loader folder.
533
    :type raddir: str
534
535
    :param radios: If we made radios this run.
536
    :type radios: bool
537
    """
538
    shutil.rmtree(osdir)
539
    if radios:
540
        shutil.rmtree(raddir)
541
542
543
def create_blitz(a_folder, swver):
544
    """
545
    Create a blitz file: a zipped archive of all app/core/radio bars.
546
547
    :param a_folder: Target folder.
548
    :type a_folder: str
549
550
    :param swver: Software version to title the blitz.
551
    :type swver: str
552
    """
553
    fname = "Blitz-{0}.zip".format(swver)
554
    with zipfile.ZipFile(fname, 'w', zipfile.ZIP_DEFLATED, allowZip64=True) as zfile:
555
        for root, dirs, files in os.walk(a_folder):
556
            del dirs
557
            for file in files:
558
                print("ZIPPING: {0}".format(utilities.stripper(file)))
559
                abs_filename = os.path.join(root, file)
560
                abs_arcname = os.path.basename(abs_filename)
561
                zfile.write(abs_filename, abs_arcname)
562
563
564
def prepends(file, pre, suf):
565
    """
566
    Check if filename starts with/ends with stuff.
567
568
    :param file: File to check.
569
    :type file: str
570
571
    :param pre: Prefix(es) to check.
572
    :type pre: str or list or tuple
573
574
    :param suf: Suffix(es) to check.
575
    :type suf: str or list or tuple
576
    """
577
    return file.startswith(pre) and file.endswith(suf)
578
579
580
def move_loaders(a_dir,
581
                 exedir_os, exedir_rad,
582
                 zipdir_os, zipdir_rad):
583
    """
584
    Move autoloaders to zipped and loaders directories in localdir.
585
586
    :param a_dir: Local directory, containing files you wish to move.
587
    :type a_dir: str
588
589
    :param exedir_os: Large autoloader .exe destination.
590
    :type exedir_os: str
591
592
    :param exedir_rad: Small autoloader .exe destination.
593
    :type exedir_rad: str
594
595
    :param zipdir_os: Large autoloader archive destination.
596
    :type zipdir_os: str
597
598
    :param zipdir_rad: Small autoloader archive destination.
599
    :type zipdir_rad: str
600
    """
601
    arx = bbconstants.ARCS
602
    pfx = bbconstants.PREFIXES
603
    loaders = [file for file in os.listdir(a_dir) if prepends(file, pfx, ".exe")]
604
    for file in loaders:
605
        print("MOVING: {0}".format(file))
606
        exedest_os = os.path.join(exedir_os, file)
607
        exedest_rad = os.path.join(exedir_rad, file)
608
        loader_sorter(file, exedest_os, exedest_rad)
609
    zippeds = [file for file in os.listdir(a_dir) if prepends(file, pfx, arx)]
610
    for file in zippeds:
611
        print("MOVING: {0}".format(file))
612
        zipdest_os = os.path.join(zipdir_os, file)
613
        zipdest_rad = os.path.join(zipdir_rad, file)
614
        loader_sorter(file, zipdest_os, zipdest_rad)
615
616
617
def loader_sorter(file, osdir, raddir):
618
    """
619
    Sort loaders based on size.
620
621
    :param file: The file to sort. Absolute paths, please.
622
    :type file: str
623
624
    :param osdir: Large file destination.
625
    :type osdir: str
626
627
    :param raddir: Small file destination.
628
    :type raddir: str
629
    """
630
    if os.path.getsize(file) > 90000000:
631
        while True:
632
            try:
633
                shutil.move(file, osdir)
634
            except shutil.Error:
635
                os.remove(file)
636
                continue
637
            break
638
    else:
639
        while True:
640
            try:
641
                shutil.move(file, raddir)
642
            except shutil.Error:
643
                os.remove(file)
644
                continue
645
            break
646
647
648
def move_bars(localdir, osdir, radiodir):
649
    """
650
    Move bar files to subfolders of a given folder.
651
652
    :param localdir: Directory to use.
653
    :type localdir: str
654
655
    :param osdir: OS file directory (large bars).
656
    :type osdir: str
657
658
    :param radiodir: Radio file directory (small bars).
659
    :type radiodir: str
660
    """
661
    for files in os.listdir(localdir):
662
        if files.endswith(".bar"):
663
            print("MOVING: {0}".format(files))
664
            bardest_os = os.path.join(osdir, files)
665
            bardest_radio = os.path.join(radiodir, files)
666
            # even the fattest radio is less than 90MB
667
            if os.path.getsize(os.path.join(localdir, files)) > 90000000:
668
                try:
669
                    shutil.move(os.path.join(localdir, files), osdir)
670
                except shutil.Error:
671
                    os.remove(bardest_os)
672
            else:
673
                try:
674
                    shutil.move(os.path.join(localdir, files), radiodir)
675
                except shutil.Error:
676
                    os.remove(bardest_radio)
677
678
679
def make_dirs(localdir, osversion, radioversion):
680
    """
681
    Create the directory tree needed for archivist/lazyloader.
682
683
    :param localdir: Root folder.
684
    :type localdir: str
685
686
    :param osversion: OS version.
687
    :type osversion: str
688
689
    :param radioversion: Radio version.
690
    :type radioversion: str
691
    """
692
    os.makedirs(localdir, exist_ok=True)
693
694
    if not os.path.exists(os.path.join(localdir, 'bars')):
695
        os.makedirs(os.path.join(localdir, 'bars'), exist_ok=True)
696
    bardir = os.path.join(localdir, 'bars')
697
    if not os.path.exists(os.path.join(bardir, osversion)):
698
        os.makedirs(os.path.join(bardir, osversion), exist_ok=True)
699
    bardir_os = os.path.join(bardir, osversion)
700
    if not os.path.exists(os.path.join(bardir, radioversion)):
701
        os.makedirs(os.path.join(bardir, radioversion), exist_ok=True)
702
    bardir_radio = os.path.join(bardir, radioversion)
703
704
    if not os.path.exists(os.path.join(localdir, 'loaders')):
705
        os.makedirs(os.path.join(localdir, 'loaders'), exist_ok=True)
706
    loaderdir = os.path.join(localdir, 'loaders')
707
    if not os.path.exists(os.path.join(loaderdir, osversion)):
708
        os.makedirs(os.path.join(loaderdir, osversion), exist_ok=True)
709
    loaderdir_os = os.path.join(loaderdir, osversion)
710
    if not os.path.exists(os.path.join(loaderdir, radioversion)):
711 View Code Duplication
        os.makedirs(os.path.join(loaderdir, radioversion), exist_ok=True)
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
712
    loaderdir_radio = os.path.join(loaderdir, radioversion)
713
714
    if not os.path.exists(os.path.join(localdir, 'zipped')):
715
        os.makedirs(os.path.join(localdir, 'zipped'), exist_ok=True)
716
    zipdir = os.path.join(localdir, 'zipped')
717
    if not os.path.exists(os.path.join(zipdir, osversion)):
718
        os.makedirs(os.path.join(zipdir, osversion), exist_ok=True)
719
    zipdir_os = os.path.join(zipdir, osversion)
720
    if not os.path.exists(os.path.join(zipdir, radioversion)):
721
        os.makedirs(os.path.join(zipdir, radioversion), exist_ok=True)
722
    zipdir_radio = os.path.join(zipdir, radioversion)
723
724
    return (bardir_os, bardir_radio, loaderdir_os, loaderdir_radio, zipdir_os, zipdir_radio)
725
726
727
def compress_config_loader(homepath=None):
728
    """
729
    Read a ConfigParser file to get compression preferences.
730
731
    :param homepath: Folder containing ini file. Default is user directory.
732
    :type homepath: str
733
    """
734 View Code Duplication
    config = configparser.ConfigParser()
0 ignored issues
show
Duplication introduced by
This code seems to be duplicated in your project.
Loading history...
735
    if homepath is None:
736
        homepath = os.path.expanduser("~")
737
    conffile = os.path.join(homepath, "bbarchivist.ini")
738
    if not os.path.exists(conffile):
739
        open(conffile, 'w').close()
740
    config.read(conffile)
741
    if not config.has_section('compression'):
742
        config['compression'] = {}
743
    compini = config['compression']
744
    method = compini.get('method', fallback="7z")
745
    if sys.version_info[1] <= 2 and method == "txz":
746
        method = "zip"  # for 3.2 compatibility
747
    return method
748
749
750
def compress_config_writer(method=None, homepath=None):
751
    """
752
    Write a ConfigParser file to store compression preferences.
753
754
    :param method: Method to use.
755
    :type method: str
756
757
    :param homepath: Folder containing ini file. Default is user directory.
758
    :type homepath: str
759
    """
760
    if method is None:
761
        method = compress_config_loader()
762
    config = configparser.ConfigParser()
763
    if homepath is None:
764
        homepath = os.path.expanduser("~")
765
    conffile = os.path.join(homepath, "bbarchivist.ini")
766
    if not os.path.exists(conffile):
767
        open(conffile, 'w').close()
768
    config.read(conffile)
769
    if not config.has_section('compression'):
770
        config['compression'] = {}
771
    config['compression']['method'] = method
772
    with open(conffile, "w") as configfile:
773
        config.write(configfile)
774