util.py 17.1 KB
Newer Older
1
2
3
# -*- coding: utf-8 -*-
# vim: ai ts=4 sts=4 et sw=4 nu

thiolliere's avatar
thiolliere committed
4
import os
rgaudin's avatar
rgaudin committed
5
6
import re
import sys
7
import json
8
9
import data
import signal
10
import base64
11
import string
rgaudin's avatar
rgaudin committed
12
import ctypes
13
14
import hashlib
import zipfile
thiolliere's avatar
util    
thiolliere committed
15
import platform
16
import tempfile
rgaudin's avatar
rgaudin committed
17
import datetime
18
import threading
rgaudin's avatar
rgaudin committed
19
import collections
20
from urllib.parse import urlparse
21

22
23
24
25
26
27
28
try:
    scandir_func = os.scandir
except AttributeError:
    import scandir

    scandir_func = scandir.scandir

29
import pytz
30
from path import Path
rgaudin's avatar
rgaudin committed
31
32
import humanfriendly

33
34
from data import cache_folder_name

35
36
37
if sys.platform == "win32":
    from win32com.shell import shellcon, shell

rgaudin's avatar
rgaudin committed
38
39
40
ONE_MiB = 2 ** 20
ONE_GiB = 2 ** 30
ONE_GB = int(1e9)
41
EXFAT_FORBIDDEN_CHARS = ["/", "\\", ":", "*", "?", '"', "<", ">", "|"]
42
PREFERENCES = None
rgaudin's avatar
rgaudin committed
43
44


45
46
47
48
49
50
51
52
53
54
STAGES = collections.OrderedDict(
    [
        ("master", "Retrieve Master Image"),
        ("download", "Download contents"),
        ("setup", "Image configuration (virtualized)"),
        ("copy", "Copy contents onto image"),
        ("move", "Post-process contents (virtualized)"),
        ("write", "SD-card creation"),
    ]
)
rgaudin's avatar
rgaudin committed
55
56
57


class ProgressHelper(object):
58
    """ progression manager for the logger
rgaudin's avatar
rgaudin committed
59
60
61
62
63
64
65
66

        must be inherited by the logger (CLILogger or Logger)

        two kinds of progresses:
            - stages (1-6) are main high-level stages
            - steps are individual tasks

        progress percentage is kept by stage (0 to 1).
67
        total percentage is calculated using current stage and its progress"""
rgaudin's avatar
rgaudin committed
68

69
70
71
72
73
74
    @staticmethod
    def _add_time(text):
        return "{time} {text}".format(
            time=datetime.datetime.now().strftime("[%X]"), text=text
        )

rgaudin's avatar
rgaudin committed
75
    def __init__(self):
76
77
78
        self.reset()

    def reset(self):
79
        self.stage_id = "init"  # id of current stage
rgaudin's avatar
rgaudin committed
80
81
82
83
84
85
86
87
        self.stage_progress = None  # percentage of current stage progress
        self.will_write = False  # wether the process will run write stage

        self.started_on = datetime.datetime.now()  # when the process started
        self.ended_on = None  # when it ended
        self.durations = {}  # records timedeltas for every ran stages

    def start(self, will_write):
88
89
        """ record logger start.
            will_write informs about whether stage 6 will take place """
rgaudin's avatar
rgaudin committed
90
91
92
93
94
95
96
        self.started_on = datetime.datetime.now()
        self.will_write = will_write

    def stop(self):
        self.ended_on = datetime.datetime.now()

    def clean_up_stage(self):
97
98
        started_on = (
            getattr(self, "stage_started_on", self.started_on)
99
            or datetime.datetime.now()
100
        )
rgaudin's avatar
rgaudin committed
101
        ended_on = datetime.datetime.now()
102
        self.durations[self.stage_id] = (started_on, ended_on, ended_on - started_on)
rgaudin's avatar
rgaudin committed
103
104
105
106
107
        self.stage_started_on = None
        self.stage_progress = None
        self.tasks = None

    def stage(self, stage_id):
108
        """ change the current stage. expects a string ID """
rgaudin's avatar
rgaudin committed
109
110
111
112
113
114
        self.clean_up_stage()  # record duration of previous stage

        self.stage_id = stage_id
        self.stage_started_on = datetime.datetime.now()
        self.update()

115
    def progress(self, numerator, denominator=1):
116
        """ record progress for the current stage """
117
118
119
120
121
122
123
124
125
126
        if numerator is not None:
            try:
                percentage = numerator / denominator
                assert percentage >= 0 and percentage <= 1
            except AssertionError:
                percentage = None
            except ZeroDivisionError:
                percentage = 1
        else:
            percentage = None
rgaudin's avatar
rgaudin committed
127
128
129
130
        self.stage_progress = percentage
        self.update()

    def ansible(self, line):
131
        """ reads logger output while in ansible
rgaudin's avatar
rgaudin committed
132
133

            detects the ansible --tasks-list call to build its list of tasks
134
            detects ansible's task calling to set progress accordingly """
rgaudin's avatar
rgaudin committed
135
136
137
138

        # display output anyway
        self.std(line)

139
        if self.stage_id not in ["setup", "move"]:
rgaudin's avatar
rgaudin committed
140
141
142
143
144
145
            return

        # detect number of task for ansiblecube phase
        if self.tasks is None and line.startswith("### TASKS ###"):
            try:
                self.tasks = [
146
147
148
                    re.search(r"^\s{6}(.*)\tTAGS\:.*$", l).groups()[0]
                    for l in line.split("^")[5:]
                ]
rgaudin's avatar
rgaudin committed
149
150
151
152
153
154
155
156
157
158
            except Exception as exp:
                print(str(exp))
                pass

        # detect tasks being executed
        # TODO: we currently record the task as complete while it just started
        if self.tasks is not None:
            # detect current task
            if line.startswith("TASK ["):
                try:
159
                    task = re.search(r"^TASK \[(.*)\] \*+$", line).groups()[0]
rgaudin's avatar
rgaudin committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
                except Exception:
                    return
                else:
                    self.step(task)
                    try:
                        task_index = self.tasks.index(task)
                    except ValueError:
                        pass
                    else:
                        self.progress(task_index / len(self.tasks))

    @property
    def stage_name(self):
        return self.get_stage_name(self.stage_id)

    @property
    def nb_of_stages(self):
        n = len(STAGES)
        return n if self.will_write else n - 1
179

rgaudin's avatar
rgaudin committed
180
181
182
    @property
    def stage_number(self):
        return self.get_stage_number(self.stage_id)
183

rgaudin's avatar
rgaudin committed
184
185
186
187
188
    @property
    def stage_numbers(self):
        return "{c}/{t}".format(c=self.stage_number, t=self.nb_of_stages)

    def get_stage_string(self, stage_id):
189
190
191
192
193
        return "[{c}/{t}] {n}".format(
            c=self.get_stage_number(stage_id),
            t=self.nb_of_stages,
            n=self.get_stage_name(stage_id),
        )
rgaudin's avatar
rgaudin committed
194
195
196
197
198
199
200
201
202
203
204
205
206

    @classmethod
    def get_stage_number(cls, stage_id):
        try:
            return list(STAGES.keys()).index(stage_id) + 1
        except Exception:
            return 0

    @classmethod
    def get_stage_name(cls, stage_id):
        return STAGES.get(stage_id, "Preparations")

    def get_overall_progress(self):
207
        """ total progression based on current stage and its progress """
rgaudin's avatar
rgaudin committed
208
209
210
211
212
213
214
215
216
217
        if not self.stage_number:
            return 0
        span = 1 / self.nb_of_stages
        lbound = span * (self.stage_number - 1)
        if self.stage_progress is None:
            return lbound
        current_progress = self.stage_progress * span
        return lbound + current_progress

    def complete(self):
218
        """ mark the logger as complete (successful) """
rgaudin's avatar
rgaudin committed
219
220
221
222
        self.clean_up_stage()
        self.stop()

    def failed(self):
223
        """ mark the logger as complete (failure) """
rgaudin's avatar
rgaudin committed
224
225
226
227
        self.clean_up_stage()
        self.stop()

    def update(self):
228
        """ update UI according to new stage/step/progress """
rgaudin's avatar
rgaudin committed
229
230
231
        raise NotImplementedError()

    def summary(self):
232
        """ textual summary of each stage's durations """
rgaudin's avatar
rgaudin committed
233
234
235
236
237
238
239
240
241
242

        # make sure we have an end datetime
        if self.ended_on is None:
            self.stop()

        self.std("*** DURATIONS SUMMARY ***")
        for stage_id in STAGES.keys():
            data = self.durations.get(stage_id)
            if data is None:
                continue
243
244
245
246
247
248
249
250
            self.std(
                "{stage}: {duration} ({start} to {end})".format(
                    stage=self.get_stage_string(stage_id),
                    duration=humanfriendly.format_timespan(data[2].total_seconds()),
                    start=data[0].strftime("%c"),
                    end=data[1].strftime("%c"),
                )
            )
rgaudin's avatar
rgaudin committed
251
        duration = self.ended_on - self.started_on
252
253
254
255
256
257
258
259
        self.std(
            "TOTAL: {duration} ({start} to {end})".format(
                duration=humanfriendly.format_timespan(duration.total_seconds()),
                start=self.started_on.strftime("%c"),
                end=self.ended_on.strftime("%c"),
            )
        )

thiolliere's avatar
util    
thiolliere committed
260

261
def get_free_space_in_dir(dirname):
thiolliere's avatar
util    
thiolliere committed
262
    """Return folder/drive free space."""
263
    if platform.system() == "Windows":
thiolliere's avatar
util    
thiolliere committed
264
        free_bytes = ctypes.c_ulonglong(0)
265
266
267
        ctypes.windll.kernel32.GetDiskFreeSpaceExW(
            ctypes.c_wchar_p(dirname), None, None, ctypes.pointer(free_bytes)
        )
thiolliere's avatar
util    
thiolliere committed
268
269
270
271
        return free_bytes.value
    else:
        st = os.statvfs(dirname)
        return st.f_bavail * st.f_frsize
thiolliere's avatar
thiolliere committed
272

273

thiolliere's avatar
thiolliere committed
274
# Thread safe class to register pid to cancel in case of abort
thiolliere's avatar
thiolliere committed
275
276
class CancelEvent:
    def __init__(self):
thiolliere's avatar
thiolliere committed
277
278
        self._lock = threading.Lock()
        self._pids = []
rgaudin's avatar
rgaudin committed
279
280
281
        self.thread = None
        self.callback = None
        self.callback_args = None
thiolliere's avatar
thiolliere committed
282
283
284
285

    def lock(self):
        return _CancelEventRegister(self._lock, self._pids)

286
    def register_thread(self, thread):
rgaudin's avatar
rgaudin committed
287
288
289
290
291
292
293
        self.thread = thread

    def unregister_thread(self):
        self.thread = None
        self.callback = None
        self.callback_args = None

thiolliere's avatar
thiolliere committed
294
    def cancel(self):
rgaudin's avatar
rgaudin committed
295
        if self.thread is not None:
296
297
            self.thread.stop()
            self.thread.join(timeout=5)  # allow proper release of handles
rgaudin's avatar
rgaudin committed
298
299
            self.unregister_thread()

thiolliere's avatar
thiolliere committed
300
301
302
303
        self._lock.acquire()
        for pid in self._pids:
            os.kill(pid, signal.SIGTERM)

304

thiolliere's avatar
thiolliere committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
class _CancelEventRegister:
    def __init__(self, lock, pids):
        self._lock = lock
        self._pids = pids

    def __enter__(self):
        self._lock.acquire()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._lock.release()

    def register(self, pid):
        if self._pids.count(pid) == 0:
            self._pids.append(pid)

    def unregister(self, pid):
        self._pids.remove(pid)
thiolliere's avatar
thiolliere committed
323

324

rgaudin's avatar
rgaudin committed
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
def human_readable_size(size, binary=True):
    if isinstance(size, (int, float)):
        num_bytes = size
    else:
        try:
            num_bytes = humanfriendly.parse_size(size)
        except Exception:
            return "NaN"
    is_neg = num_bytes < 0
    if is_neg:
        num_bytes = abs(num_bytes)
    output = humanfriendly.format_size(num_bytes, binary=binary)
    if is_neg:
        return "- {}".format(output)
    return output
thiolliere's avatar
thiolliere committed
340

341

rgaudin's avatar
rgaudin committed
342
343
344
345
class CLILogger(ProgressHelper):
    def __init__(self):
        super(CLILogger, self).__init__()

346
347
348
349
350
351
352
    @property
    def on_tty(self):
        return True

    def flash(self, line):
        self.raw_std(line + "\r")

353
354
    def step(self, step, end=None):
        self.p("--> {}".format(step), color="34", end=end)
rgaudin's avatar
rgaudin committed
355

356
357
    def err(self, err, end=None):
        self.p(err, color="31", end=end)
rgaudin's avatar
rgaudin committed
358

359
360
    def succ(self, succ, end=None):
        self.p(succ, color="32", end=end)
rgaudin's avatar
rgaudin committed
361

rgaudin's avatar
rgaudin committed
362
363
364
365
366
367
368
369
    def raw_std(self, std):
        sys.stdout.write(std)

    def std(self, std, end=None):
        self.p(std, end=end, flush=True)

    def p(self, text, color=None, end=None, flush=False):
        if color is not None and sys.platform != 32:
370
371
372
            text = "\033[00;{col}m{text}\033[00m".format(
                col=color, text=self._add_time(text)
            )
rgaudin's avatar
rgaudin committed
373
374
375
376
        print(text, end=end, flush=flush)

    def complete(self):
        super(CLILogger, self).complete()
rgaudin's avatar
rgaudin committed
377
        self.succ("Installation succeded.")
rgaudin's avatar
rgaudin committed
378
379
380
381
382
383

    def failed(self, error="?"):
        super(CLILogger, self).failed()
        self.err("Installation failed: {}".format(error))

    def update(self):
384
385
386
387
388
389
390
391
392
        self.p(
            "[STAGE {nums}: {name} - {pc:.0f}%]".format(
                nums=self.stage_numbers,
                name=self.stage_name,
                pc=self.get_overall_progress() * 100,
            ),
            color="35",
        )

rgaudin's avatar
rgaudin committed
393
394
395
396
397
398
399
400

def get_checksum(fpath, func=hashlib.sha256):
    h = func()
    with open(fpath, "rb") as f:
        for chunk in iter(lambda: f.read(ONE_MiB * 8), b""):
            h.update(chunk)
    return h.hexdigest()

401

rgaudin's avatar
rgaudin committed
402
def get_cache(build_folder):
403
    fpath = os.path.join(build_folder, cache_folder_name)
rgaudin's avatar
rgaudin committed
404
405
406
    os.makedirs(fpath, exist_ok=True)
    return fpath

407

rgaudin's avatar
rgaudin committed
408
409
410
def get_temp_folder(in_path):
    os.makedirs(in_path, exist_ok=True)
    return tempfile.mkdtemp(dir=in_path)
411

412

413
def relpathto(dest):
414
    """ relative path to an absolute one """
415
416
417
    if dest is None:
        return None
    return str(Path(dest).relpath())
418

419

420
def b64encode(fpath):
421
    """ base64 string of a binary file """
422
    with open(fpath, "rb") as fp:
423
424
        return base64.b64encode(fp.read()).decode("utf-8")

425
426

def b64decode(fname, data, to):
427
    """ write back a binary file from its fname and base64 string """
428
    fpath = os.path.join(to, fname)
429
    with open(fpath, "wb") as fp:
430
431
        fp.write(base64.b64decode(data))
    return fpath
432
433
434


def exfat_fnames_filter(fname):
435
    """ whether supplied fname is valid exfat fname or not """
436
437
438
439
440
    # TODO: check for chars U+0000 to U+001F
    return sum([1 for x in EXFAT_FORBIDDEN_CHARS if x in fname]) == 0


def ensure_zip_exfat_compatible(fpath):
441
    """ wether supplied ZIP archive at fpath contains exfat-OK file names
442

443
        boolean, [erroneous, file, names]"""
444
445
    bad_fnames = []
    try:
446
        with zipfile.ZipFile(fpath, "r") as zipf:
447
448
449
450
            # loop over all file names in the ZIP
            for path in zipf.namelist():
                # loop over all parts (folder, subfolder(s), fname)
                for part in Path(path).splitall():
451
                    if not exfat_fnames_filter(part) and part not in bad_fnames:
452
453
454
455
                        bad_fnames.append(part)
    except Exception as exp:
        return False, [str(exp)]
    return len(bad_fnames) == 0, bad_fnames
456
457


458
459
460
def check_user_inputs(
    project_name, language, timezone, admin_login, admin_pwd, wifi_pwd=None
):
461

462
463
464
465
466
467
    allowed_chars = set(
        string.ascii_uppercase + string.ascii_lowercase + string.digits + "-" + " "
    )
    valid_project_name = (
        len(project_name) >= 1
        and len(project_name) <= 64
468
        and set(project_name) <= allowed_chars
469
    )
470

rgaudin's avatar
rgaudin committed
471
    valid_language = language in dict(data.hotspot_languages).keys()
472
473
474

    valid_timezone = timezone in pytz.common_timezones

475
    valid_wifi_pwd = (
rgaudin's avatar
rgaudin committed
476
477
        len(wifi_pwd) >= 8
        and len(wifi_pwd) <= 31
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
        and set(wifi_pwd)
        <= set(
            string.ascii_uppercase + string.ascii_lowercase + string.digits + "-" + "_"
        )
        if wifi_pwd is not None
        else True
    )

    valid_admin_login = len(admin_login) <= 31 and set(admin_login) <= set(
        string.ascii_uppercase + string.ascii_lowercase + string.digits + "-" + "_"
    )
    valid_admin_pwd = len(admin_pwd) <= 31 and set(admin_pwd) <= set(
        string.ascii_uppercase + string.ascii_lowercase + string.digits + "-" + "_"
    )

    return (
        valid_project_name,
        valid_language,
        valid_timezone,
        valid_wifi_pwd,
        valid_admin_login,
        valid_admin_pwd,
    )
rgaudin's avatar
rgaudin committed
501
502
503


def get_adjusted_image_size(size):
504
    """ save some space to accomodate real SD card sizes
rgaudin's avatar
rgaudin committed
505

506
        the larger the SD card, the larger the loss space is """
rgaudin's avatar
rgaudin committed
507

508
509
510
511
512
    # if size is not a rounded GB multiple, assume it's OK
    if not size % ONE_GB == 0:
        return size

    rate = .97 if size / ONE_GB <= 16 else .96
rgaudin's avatar
rgaudin committed
513
    return int(size * rate)
514
515
516
517
518
519
520
521
522
523
524


def get_folder_size(path):
    """ total size in bytes of a folder (recurse file stat) """
    total = 0
    for entry in scandir_func(path):
        if entry.is_file():
            total += os.path.getsize(entry.path)
        elif entry.is_dir():
            total += get_folder_size(entry.path)
    return total
525
526
527
528
529
530
531


def split_proxy(proxy_url):
    """ return a (username:password@host, port) tuple from proxy URL """
    try:
        url = urlparse(proxy_url)
        if url.port:
532
533
534
535
536
            netloc = url.netloc[: -(len(str(url.port)) + 1)]  # remove port from netloc
            port = str(url.port)
        else:
            netloc = url.netloc
            port = ""
537
    except Exception:
538
        # could not parse URL, attempt something
539
        parts = proxy_url.rsplit(":", 1)
540
541
542
        netloc = parts[0].split("://", 1)[-1]
        port = parts[1] if len(parts) > 1 else ""
    return netloc, port
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588


def get_prefs_path():
    """ full path to our preferences JSON file """
    fname = "kiwix-hotspot.prefs"
    homedir = os.path.expanduser("~")
    if sys.platform == "win32":
        try:
            prefsdir = shell.SHGetFolderPath(0, shellcon.CSIDL_LOCAL_APPDATA, 0, 0)
        except Exception:
            prefsdir = homedir
    elif sys.platform == "linux":
        prefsdir = os.path.join(homedir, ".config")
    elif sys.platform == "darwin":
        prefsdir = os.path.join(homedir, "Library", "Application Support")
    else:
        raise NotImplementedError("OS Not Supported")

    os.makedirs(prefsdir, exist_ok=True)  # create folder should it not exist
    return os.path.join(prefsdir, fname)


def read_preferences():
    """ read and parse JSON preference file """
    fpath = get_prefs_path()
    if os.path.exists(fpath):
        try:
            with open(fpath, "r") as fd:
                return json.load(fd)
        except Exception as exp:
            print("Failed to parse prefs at {}: {}".format(fpath, exp))
    return {}


def get_prefs(force_reload=False):
    """ cached-shortcut to PREFERENCES """
    global PREFERENCES
    if PREFERENCES is None or force_reload:
        PREFERENCES = read_preferences()
    return PREFERENCES


def save_prefs(prefs, auto_reload=True):
    """ save a passed preferences dict to the preferences file on disk """
    fpath = get_prefs_path()
    try:
rgaudin's avatar
rgaudin committed
589
        with open(fpath, "w", encoding="utf-8") as fd:
590
591
592
593
594
595
596
597
            json.dump(prefs, fd, indent=4)
    except Exception as exp:
        print("Failed to save prefs to {}: {}".format(fpath, exp))
        return False

    if auto_reload:
        get_prefs(force_reload=True)
    return True