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

rgaudin's avatar
rgaudin committed
4
import os
5
import re
rgaudin's avatar
rgaudin committed
6
import sys
7
import time
8
import shlex
9
import signal
rgaudin's avatar
rgaudin committed
10
import ctypes
11
import tempfile
rgaudin's avatar
rgaudin committed
12
import threading
thiolliere's avatar
thiolliere committed
13
14
import subprocess

15
import data
rgaudin's avatar
rgaudin committed
16
from util import CLILogger
rgaudin's avatar
rgaudin committed
17
18


19
20
21
22
23
24
# windows-only flags to prevent sleep on executing thread
WINDOWS_SLEEP_FLAGS = {
    # Enables away mode. This value must be specified with ES_CONTINUOUS.
    # Away mode should be used only by media-recording and media-distribution
    # applications that must perform critical background processing
    # on desktop computers while the computer appears to be sleeping.
25
    "ES_AWAYMODE_REQUIRED": 0x00000040,
26
27
    # Informs the system that the state being set should remain in effect until
    # the next call that uses ES_CONTINUOUS and one of the other state flags is cleared.
28
    "ES_CONTINUOUS": 0x80000000,
29
    # Forces the display to be on by resetting the display idle timer.
30
    "ES_DISPLAY_REQUIRED": 0x00000002,
31
    # Forces the system to be in the working state by resetting the system idle timer.
32
    "ES_SYSTEM_REQUIRED": 0x00000001,
33
34
35
}


thiolliere's avatar
thiolliere committed
36
37
38
39
class CheckCallException(Exception):
    def __init__(self, msg):
        Exception(self, msg)

rgaudin's avatar
rgaudin committed
40

thiolliere's avatar
thiolliere committed
41
def startup_info_args():
42
    if hasattr(subprocess, "STARTUPINFO"):
thiolliere's avatar
thiolliere committed
43
44
45
46
47
        # On Windows, subprocess calls will pop up a command window by default
        # when run from Pyinstaller with the ``--noconsole`` option. Avoid this
        # distraction.
        si = subprocess.STARTUPINFO()
        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
48
        cf = subprocess.CREATE_NEW_PROCESS_GROUP
thiolliere's avatar
thiolliere committed
49
    else:
rgaudin's avatar
rgaudin committed
50
51
        si = None
        cf = 0
52
    return {"startupinfo": si, "creationflags": cf}
thiolliere's avatar
thiolliere committed
53

rgaudin's avatar
rgaudin committed
54

55
56
57
58
def subprocess_pretty_call(
    cmd, logger, stdin=None, check=False, decode=False, as_admin=False
):
    """ flexible subprocess helper running separately and using the logger
rgaudin's avatar
rgaudin committed
59
60
61
62
63
64

        cmd: the command to be run
        logger: the logger to send debug output to
        stdin: pipe input into the command
        check: whether it should raise on non-zero return code
        decode: whether it should decode output (bytes) into UTF-8 str
65
        as_admin: whether the command should be run as root/admin """
rgaudin's avatar
rgaudin committed
66
67
68
69
70

    if as_admin:
        if sys.platform == "win32":
            if logger is not None:
                logger.std("Call (as admin): " + str(cmd))
71
            return run_as_win_admin(cmd, logger)
rgaudin's avatar
rgaudin committed
72
73

        from_cli = logger is None or type(logger) == CLILogger
74
        cmd = get_admin_command(cmd, from_gui=not from_cli, logger=logger)
rgaudin's avatar
rgaudin committed
75

thiolliere's avatar
thiolliere committed
76
    # We should use subprocess.run but it is not available in python3.4
77
78
79
80
81
82
83
    process = subprocess.Popen(
        cmd,
        stdin=stdin,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        **startup_info_args()
    )
rgaudin's avatar
rgaudin committed
84
85
86
87

    if logger is not None:
        logger.std("Call: " + str(process.args))

thiolliere's avatar
thiolliere committed
88
89
    process.wait()

90
91
92
93
94
    lines = (
        [l.decode("utf-8", "ignore") for l in process.stdout.readlines()]
        if decode
        else process.stdout.readlines()
    )
rgaudin's avatar
rgaudin committed
95
96
97
98
99
100
101
102
103
104
105
106
107
108

    if logger is not None:
        for line in lines:
            logger.raw_std(line if decode else line.decode("utf-8", "ignore"))

    if check:
        if process.returncode != 0:
            raise CheckCallException("Process %s failed" % process.args)
        return lines

    return process.returncode, lines


def subprocess_pretty_check_call(cmd, logger, stdin=None, as_admin=False):
109
110
111
    return subprocess_pretty_call(
        cmd=cmd, logger=logger, stdin=stdin, check=True, as_admin=as_admin
    )
rgaudin's avatar
rgaudin committed
112
113


rgaudin's avatar
rgaudin committed
114
115
116
117
118
119
120
def subprocess_timed_output(cmd, logger, timeout=10):
    logger.std("Getting output of " + str(cmd))
    return subprocess.check_output(
        cmd, universal_newlines=True, timeout=timeout
    ).splitlines()


121
def subprocess_external(cmd, logger):
122
    """ spawn a new process without capturing nor watching it """
123
124
125
126
    logger.std("Opening: " + str(cmd))
    subprocess.Popen(cmd)


rgaudin's avatar
rgaudin committed
127
def is_admin():
128
    """ whether current process is ran as Windows Admin or unix root """
rgaudin's avatar
rgaudin committed
129
130
131
132
133
134
135
136
137
    if sys.platform == "win32":
        try:
            return ctypes.windll.shell32.IsUserAnAdmin()
        except Exception:
            return False
    return os.getuid() == 0


def run_as_win_admin(command, logger):
138
    """ run specified command with admin rights """
139
    params = " ".join(['"{}"'.format(x) for x in command[1:]]).strip()
140
    rc = ctypes.windll.shell32.ShellExecuteW(None, "runas", command[0], params, None, 1)
rgaudin's avatar
rgaudin committed
141
142
143
144
145
146
    # ShellExecuteW returns 5 if user chose not to elevate
    if rc == 5:
        raise PermissionError()
    return rc


147
def get_admin_command(command, from_gui, logger, log_to=None):
148
    """ updated command to run it as root on macos or linux
thiolliere's avatar
thiolliere committed
149

150
        from_gui: whether called via GUI. Using cli sudo if not """
thiolliere's avatar
thiolliere committed
151

rgaudin's avatar
rgaudin committed
152
153
    if not from_gui:
        return ["sudo"] + command
thiolliere's avatar
thiolliere committed
154

rgaudin's avatar
rgaudin committed
155
    if sys.platform == "darwin":
156
157
158
159
        # write command to a separate temp bash script
        script = (
            "#!/bin/bash\n\n{command} 2>&1 {redir}\n\n"
            'if [ $? -eq 1 ]; then\n    echo "!!! echer returned 1" {redir}\n'
160
            "    exit 11\nfi\n\n".format(
161
162
163
164
                command=" ".join([shlex.quote(cmd) for cmd in command]),
                redir=">>{}".format(log_to) if log_to else "",
            )
        )
165
166
167
168

        # add script content to logger
        logger.raw_std(script)

169
170
171
172
173
174
175
176
177
178
        with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as fd:
            fd.write(script)
            fd.seek(0)

            return [
                "/usr/bin/osascript",
                "-e",
                'do shell script "/bin/bash {command}" '
                "with administrator privileges".format(command=fd.name),
            ]
rgaudin's avatar
rgaudin committed
179
180
    if sys.platform == "linux":
        return ["pkexec"] + command
rgaudin's avatar
rgaudin committed
181
182


183
class EtcherWriterThread(threading.Thread):
184
185
186
187
188
189
190
191
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._should_stop = False  # stop flag
        self.exp = None  # exception to be re-raised by caller

    def stop(self):
        self._should_stop = True

192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
    @classmethod
    def show_log(cls, logger, log_to_file, log_file, process, eof=False):
        if log_to_file:
            try:
                with open(log_file.name, "r") as f:
                    lines = f.readlines()
                    if len(lines) >= 2:
                        lines.pop()
                    # working
                    if "Validating" in lines[-1] or "Flashing" in lines[-1]:
                        logger.std(lines[-1].replace("\x1b[1A", "").strip())
                    elif "[1A" in lines[-1]:  # still working but between progress
                        logger.std(lines[-2].replace("\x1b[1A", "").strip())
                    else:  # probably at end of file
                        for line in lines[-5:]:
                            logger.std(line.replace("\x1b[1A", "").strip())
            except Exception as exp:
                logger.err("Failed to read etcher log output: {}".format(exp))

        if not log_to_file or eof:
            for line in process.stdout:
                logger.raw_std(line.decode("utf-8", "ignore"))

215
    def run(self,):
216
217
        image_fpath, device_fpath, logger = self._args

218
        logger.step("Copy image to sd card using etcher-cli")
219

220
        from_cli = logger is None or type(logger) == CLILogger
221
        cmd, log_to_file, log_file = get_etcher_command(
222
            image_fpath, device_fpath, logger, from_cli
223
        )
rgaudin's avatar
rgaudin committed
224

225
226
227
228
        process = subprocess.Popen(
            cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, **startup_info_args()
        )
        logger.std("Starting Etcher: " + str(process.args))
229

rgaudin's avatar
rgaudin committed
230
231
232
233
234
        # intervals in second
        sleep_interval = 2
        log_interval = 60

        counter = 0
235
        while process.poll() is None:
rgaudin's avatar
rgaudin committed
236
            counter += 1
237
238
            if self._should_stop:  # on cancel
                logger.std(". cancelling...")
239
240
                break

rgaudin's avatar
rgaudin committed
241
242
243
244
245
246
247
248
            time.sleep(sleep_interval)
            # increment sleep counter until we reach log interval
            if counter < log_interval // sleep_interval:
                counter += 1
                continue

            # reset counter and display log
            counter = 0
249
            self.show_log(logger, log_to_file, log_file, process)
250

251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
        try:
            logger.std(". has process exited?")
            process.wait(timeout=2)
        except subprocess.TimeoutExpired:
            logger.std(". process exited")
            # send ctrl^c
            if sys.platform == "win32":
                logger.std(". sending ctrl^C")
                process.send_signal(signal.CTRL_C_EVENT)
                process.send_signal(signal.CTRL_BREAK_EVENT)
                time.sleep(2)
            if process.poll() is None:
                logger.std(". sending SIGTERM")
                process.terminate()  # send SIGTERM
                time.sleep(2)
            if process.poll() is None:
                logger.std(". sending SIGKILL")
                process.kill()  # send SIGKILL (SIGTERM again on windows)
                time.sleep(2)
        else:
            logger.std(". process exited")
            if not process.returncode == 0:
                self.exp = CheckCallException(
                    "Process returned {}".format(process.returncode)
                )
276
277
278
279
280
281
282
283
284
285
286

        # capture last output
        self.show_log(logger, log_to_file, log_file, process, eof=True)

        if log_to_file:
            log_file.close()
            try:
                os.unlink(log_file.name)
            except Exception as exp:
                logger.err(str(exp))

287
        logger.std(". process done")
288
        logger.progress(1)
289
290
291
292
293
294


def prevent_sleep(logger):
    if sys.platform == "win32":
        logger.std("Setting ES_SYSTEM_REQUIRED mode to current thread")
        ctypes.windll.kernel32.SetThreadExecutionState(
295
296
297
298
            WINDOWS_SLEEP_FLAGS["ES_CONTINUOUS"]
            | WINDOWS_SLEEP_FLAGS["ES_SYSTEM_REQUIRED"]
            | WINDOWS_SLEEP_FLAGS["ES_DISPLAY_REQUIRED"]
        )
299
300
301
        return

    if sys.platform == "linux":
302

303
304
        def make_unmapped_window(wm_name):
            from Xlib import display
305

306
307
308
309
310
311
312
313
314
315
316
317
            screen = display.Display().screen()
            window = screen.root.create_window(0, 0, 1, 1, 0, screen.root_depth)
            window.set_wm_name(wm_name)
            window.set_wm_protocols([])
            return window

        logger.std("Suspending xdg-screensaver")
        wid = None
        try:
            # Create window to use with xdg-screensaver
            window = make_unmapped_window("caffeinate")
            wid = hex(window.id)
318
            cmd = ["/usr/bin/xdg-screensaver", "suspend", wid]
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
            logger.std("Calling {}".format(cmd))
            p = subprocess.Popen(" ".join(cmd), shell=True)
            p.wait()
            if not p.returncode == 0:
                raise OSError("xdg-screensaver returned {}".format(p.returncode))
        except Exception as exp:
            logger.err("Unable to disable sleep. Please do it manually.")
        return wid

    if sys.platform == "darwin":
        cmd = ["/usr/bin/caffeinate", "-dsi"]
        logger.std("Calling {}".format(cmd))
        process = subprocess.Popen(cmd, **startup_info_args())

        return process

335

336
337
338
339
def restore_sleep_policy(reference, logger):
    if sys.platform == "win32":
        logger.std("Restoring ES_CONTINUOUS mode to current thread")
        ctypes.windll.kernel32.SetThreadExecutionState(
340
341
            WINDOWS_SLEEP_FLAGS["ES_CONTINUOUS"]
        )
342
343
344
345
346
347
        return

    if sys.platform == "linux":
        logger.std("Resuming xdg-screensaver (wid #{})".format(reference))
        if reference is not None:
            subprocess_pretty_call(
348
349
                ["/usr/bin/xdg-screensaver", "resume", reference], logger
            )
350
351
352
353
354
355
356
        return

    if sys.platform == "darwin":
        logger.std("Stopping caffeinate process #{}".format(reference.pid))
        reference.kill()
        reference.wait(5)
        return
357
358


359
def get_etcher_command(image_fpath, device_fpath, logger, from_cli):
360
361
362
363

    # on macOS, GUI sudo captures stdout so we use a log file
    log_to_file = not from_cli and sys.platform == "darwin"
    if log_to_file:
rgaudin's avatar
rgaudin committed
364
365
366
        log_file = tempfile.NamedTemporaryFile(
            suffix=".log", delete=False, encoding="utf-8"
        )
367
368
369
370
    else:
        log_file = None

    cmd = [
371
372
373
374
375
        os.path.join(
            data.data_dir,
            "etcher-cli",
            "etcher" if sys.platform == "win32" else "balena-etcher",
        ),
376
377
378
379
380
381
382
        "-c",
        "-y",
        "-u",
        "-d",
        device_fpath,
        image_fpath,
    ]
383

384
385
386
    # handle sudo or GUI alternative for linux and macOS
    if sys.platform in ("linux", "darwin"):
        cmd = get_admin_command(
387
388
389
390
            cmd,
            from_gui=not from_cli,
            logger=logger,
            log_to=log_file.name if log_to_file else None,
391
392
393
394
395
        )

    return cmd, log_to_file, log_file


396
def flash_image_with_etcher(image_fpath, device_fpath, retcode, from_cli=False):
397
398
399
400
401
402
403
    """ flash an image onto SD-card

        use only with small image as there is no output capture on OSX
        and it is not really cancellable.
        retcode is a multiprocessing.Value """
    logger = CLILogger()
    cmd, log_to_file, log_file = get_etcher_command(
404
        image_fpath, device_fpath, logger, from_cli
405
406
407
    )
    returncode, _ = subprocess_pretty_call(cmd, check=False, logger=logger)
    retcode.value = returncode
408
409
    if log_to_file:
        try:
410
411
            subprocess_pretty_call(["/bin/cat", log_file.name], logger, decode=True)
            log_file.close()
412
413
414
            os.unlink(log_file.name)
        except Exception as exp:
            logger.err(str(exp))
415
416
417
418
419
420
421
422
    return returncode == 0


def sd_has_single_partition(sd_card, logger):
    """ whether sd_card consists of a single partition (expected to be clean) """
    try:
        if sys.platform == "darwin":
            disk_prefix = re.sub(r"\/dev\/disk([0-9]+)", r"disk\1s", sd_card)
rgaudin's avatar
rgaudin committed
423
            lines = subprocess_timed_output(["diskutil", "list", sd_card], logger)
424
425
426
427
428
429
430
431
432
433
434
435
            nb_partitions = len(
                [
                    line.strip().rsplit(" ", 1)[-1].replace(disk_prefix, "").strip()
                    for line in lines
                    if disk_prefix in line
                ]
            )
            return nb_partitions == 1
        elif sys.platform == "win32":
            disk_prefix = re.sub(
                r".+PHYSICALDRIVE([0-9+])", r"Disk #\1, Partition #", sd_card
            )
rgaudin's avatar
rgaudin committed
436
            lines = subprocess_timed_output(["wmic", "partition"], logger)
437
438
439
440
441
442
443
444
445
446
            nb_partitions = len(
                [
                    re.sub(r".+" + disk_prefix + r"([0-9]+).+", r"\1", line)
                    for line in lines
                    if disk_prefix in line
                ]
            )
            return nb_partitions == 1
        elif sys.platform == "linux":
            disk_prefix = re.sub(r"\/dev\/([a-z0-9]+)", r"─\1", sd_card)
rgaudin's avatar
rgaudin committed
447
            lines = subprocess_timed_output(["/bin/lsblk", sd_card], logger)
448
449
450
451
452
453
454
455
456
457
458
            nb_partitions = len(
                [
                    line.strip().split(" ", 1)[0].replace(disk_prefix, "").strip()
                    for line in lines
                    if disk_prefix in line
                ]
            )
            return nb_partitions == 1
    except Exception as exp:
        logger.err(str(exp))
        return False