fslinstaller.py 62.6 KB
Newer Older
1
#!/usr/bin/env python
2
3
4
#
# FSL installer script.
#
Paul McCarthy's avatar
Paul McCarthy committed
5
6
7
8
"""This is the FSL installation script. It can be used to install FSL, or
to update an existing FSL installation.  This script can be executed with
Python 2.7 or newer.
"""
9
10


11
from __future__ import print_function, division, unicode_literals
12

13
14
15
import functools      as ft
import os.path        as op
import subprocess     as sp
16
import textwrap       as tw
17
18
19
import                   argparse
import                   contextlib
import                   getpass
Paul McCarthy's avatar
Paul McCarthy committed
20
import                   hashlib
21
22
import                   json
import                   logging
23
import                   os
24
import                   platform
25
import                   re
26
import                   readline
27
28
import                   shlex
import                   shutil
29
import                   sys
30
import                   tempfile
31
32
import                   threading
import                   time
Paul McCarthy's avatar
Paul McCarthy committed
33
import                   traceback
34

35
# TODO check py2/3
36
37
38
try:
    import urllib.request as urlrequest
except ImportError:
39
    import urllib as urlrequest
40

41
42
43
44
45
46
try:                import queue
except ImportError: import Queue as queue


PY2 = sys.version[0] == '2'

47

48
49
50
log = logging.getLogger(__name__)


Paul McCarthy's avatar
Paul McCarthy committed
51
52
# this sometimes gets set to fslinstaller.pyc, so rstrip c
__absfile__ = op.abspath(__file__).rstrip('c')
53
54


Paul McCarthy's avatar
Paul McCarthy committed
55
__version__ = '1.3.0'
Paul McCarthy's avatar
Paul McCarthy committed
56
57
"""Installer script version number. This is automatically updated
whenever a new version of the installer script is released.
Paul McCarthy's avatar
Paul McCarthy committed
58
59
60
61
62
63
64
"""


DEFAULT_INSTALLATION_DIRECTORY = '/usr/local/fsl'
"""Default FSL installation directory. """


Paul McCarthy's avatar
Paul McCarthy committed
65
FSL_INSTALLER_MANIFEST = 'http://18.133.213.73/releases/manifest.json'
Paul McCarthy's avatar
Paul McCarthy committed
66
"""URL to download the FSL installer manifest file from. The installer
67
68
69
70
71
manifest file is a JSON file which contains information about available FSL
versions.

See the Context.download_manifest function, and an example manifest file
in test/data/manifest.json, for more details.
Paul McCarthy's avatar
Paul McCarthy committed
72

Paul McCarthy's avatar
Paul McCarthy committed
73
74
A custom manifest URL can be specified with the -a/--manifest command-line
option.
75
76
77
"""


78
79
FIRST_FSL_CONDA_RELEASE = '6.0.6'
"""Oldest conda-based FSL version that can be updated in-place by this
Paul McCarthy's avatar
Paul McCarthy committed
80
installer script. Versions older than this will need to be overwritten.
81
82
83
84
85
86
"""


@ft.total_ordering
class Version(object):
    """Class to represent and compare version strings.  Accepted version
Paul McCarthy's avatar
Paul McCarthy committed
87
    strings are of the form W.X.Y.Z, where W, X, Y, and Z are all integers.
88
89
    """
    def __init__(self, verstr):
90
91
92
93
94
95
96
97
98
99
100
101
        # Version identifiers for official FSL
        # releases will have up to four
        # components (X.Y.Z.W), but We accept
        # any number of (integer) components,
        # as internal releases may have more.
        components = []

        for comp in verstr.split('.'):
            try:              components.append(int(comp))
            except Exception: break

        self.components = components
Paul McCarthy's avatar
Paul McCarthy committed
102
        self.verstr     = verstr
103
104
105
106
107

    def __str__(self):
        return self.verstr

    def __eq__(self, other):
Paul McCarthy's avatar
Paul McCarthy committed
108
109
110
111
        for sn, on in zip(self.components, other.components):
            if sn != on:
                 return False
        return len(self.components) == len(other.components)
112
113

    def __lt__(self, other):
Paul McCarthy's avatar
Paul McCarthy committed
114
        for p1, p2 in zip(self.components, other.components):
115
116
            if p1 < p2: return True
            if p1 > p2: return False
117
        return len(self.components) < len(other.components)
118
119


Paul McCarthy's avatar
Paul McCarthy committed
120
121
122
class Context(object):
    """Bag of information and settings created in main, and passed around
    this script.
123
124
125

    Several settings are lazily evaluated on first access, but once evaluated,
    their values are immutable.
Paul McCarthy's avatar
Paul McCarthy committed
126
127
128
    """

    def __init__(self, args):
129
130
131
        """Create the context with the argparse.Namespace object containing
        parsed command-line arguments.
        """
132

133
        self.args  = args
Paul McCarthy's avatar
Paul McCarthy committed
134
        self.shell = op.basename(os.environ.get('SHELL', 'sh')).lower()
Paul McCarthy's avatar
Paul McCarthy committed
135
136

        # These attributes are updated on-demand via
137
138
139
140
141
142
143
        # the property accessors defined below, or all
        # all updated via the finalise-settings method.
        self.__platform       = None
        self.__cuda           = None
        self.__manifest       = None
        self.__build          = None
        self.__destdir        = None
Paul McCarthy's avatar
Paul McCarthy committed
144
145
        self.__need_admin     = None
        self.__admin_password = None
146

Paul McCarthy's avatar
Paul McCarthy committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
        # These attributes are set by main - exists is
        # a flag denoting whether the dest dir already
        # exists, and update is the version string of
        # the existing FSL installation if the user
        # has selected to update it, or None otherwise.
        self.exists = False
        self.update = None

        # If the destination directory already exists,
        # and the user chooses to overwrite it, it is
        # moved so that, if the installation fails, it
        # can be restored. The new path is stored
        # here - refer to overwrite_destdir.
        self.old_destdir = None

162
        # The download_fsl_environment function stores
163
164
165
        # the path to the FSL conda environment file,
        # list of conda channels, and fsl-base version,
        # here.
166
167
        self.environment_file     = None
        self.environment_channels = None
168
        self.fsl_base_version     = None
Paul McCarthy's avatar
Paul McCarthy committed
169
170
171
172
173

        # The config_logging function stores the path
        # to the fslinstaller log file here.
        self.logfile = None

174
175

    def finalise_settings(self):
176
177
        """Finalise values for all information and settings in the Context.
        """
Paul McCarthy's avatar
Paul McCarthy committed
178
        self.manifest
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
        self.platform
        self.cuda
        self.build
        self.destdir
        self.need_admin
        self.admin_password


    @property
    def platform(self):
        """The platform we are running on, e.g. "linux-64", "macos-64". """
        if self.__platform is None:
            self.__platform = Context.identify_platform()
        return self.__platform


    @property
    def cuda(self):
        """The available CUDA version, or a CUDA version requested by the user.
        """
        if self.__cuda is not None:
            return self.__cuda
        if self.args.cuda is not None:
            self.__cuda = self.args.cuda
        if self.__cuda is None:
            self.__cuda = Context.identify_cuda()
        return self.__cuda


    @property
    def build(self):
        """Returns a suitable FSL build (a dictionary entry from the FSL
211
212
        installer manifest) for the target platform and requested FSL/CUDA
        versions.
213
214
215
216
217
218
219
220
221
222

        The returned dictionary has the following elements:
          - 'version'      FSL version.
          - 'platform':    Platform identifier (e.g. 'linux-64')
          - 'environment': Environment file to download
          - 'sha256':      Checksum of environment file
          - 'output':      Number of lines of expected output, for reporting
                           progress
          - 'cuda':        X.Y CUDA version, if a CUDA-enabled version of FSL
                           is to be installed.
223
224
        """

Paul McCarthy's avatar
Paul McCarthy committed
225
226
227
228
229
        if self.__build is not None:
            return self.__build

        # defaults to "latest" if
        # not specified by the user
230
231
232
        fslversion = self.args.fslversion

        if fslversion not in self.manifest['versions']:
Paul McCarthy's avatar
Paul McCarthy committed
233
234
            available = ', '.join(self.manifest['versions'].keys())
            raise Exception(
Paul McCarthy's avatar
Paul McCarthy committed
235
                'FSL version "{}" is not available - available '
Paul McCarthy's avatar
Paul McCarthy committed
236
                'versions: {}'.format(fslversion, available))
237
238
239
240

        if fslversion == 'latest':
            fslversion = self.manifest['versions']['latest']

Paul McCarthy's avatar
Paul McCarthy committed
241
242
243
244
245
246
247
        # Find refs to all compatible builds,
        # separating the default (no CUDA) build
        # from CUDA-enabled builds. We assume
        # that there is only one default build
        # for each platform.
        default    = None
        candidates = []
248
249

        for build in self.manifest['versions'][fslversion]:
Paul McCarthy's avatar
Paul McCarthy committed
250
251
252
253
254
255
256
257
            if build['platform'] == self.platform:
                if build.get('cuda', None) is None:
                    default = build
                else:
                    candidates.append(build)

        if (default is None) and (len(candidates) == 0):
            raise Exception(
258
259
260
                'Cannot find a version of FSL matching platform '
                '{} and CUDA {}'.format(self.platform, self.cuda))

Paul McCarthy's avatar
Paul McCarthy committed
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        # If we have CUDA (or the user has
        # specifically requested a CUDA build),
        # try and find a suitable build
        match = default
        if self.cuda is not None:
            candidates = sorted(candidates, key=lambda b: float(b['cuda']))

            for build in reversed(candidates):
                if self.cuda >= float(build['cuda']):
                    match = build
                    break
            else:
                available = [b['cuda'] for b in candidates]
                printmsg('Could not find a suitable FSL CUDA '
                         'build for CUDA version {} (available: '
                         '{}. Installing default (non-CUDA) '
                         'FSL build.'.format(self.cuda, available),
                         WARNING)
                printmsg('You can use the --cuda command-line option '
                         'to install a FSL build that is compatible '
                         'with a specific CUDA version', INFO)

        printmsg('FSL {} [CUDA: {}] selected for installation'.format(
            match['version'], match.get('cuda', 'n/a')))

        self.__build = match
287
        return match
Paul McCarthy's avatar
Paul McCarthy committed
288
289
290
291


    @property
    def destdir(self):
292
293
294
        """Installation directory. If not specified at the command line, the
        user is prompted to enter a directory.
        """
295
296
297
298
299
300
301
302
303

        if self.__destdir is not None:
            return self.__destdir

        # The loop below validates the destination directory
        # both when specified at commmand line or
        # interactively.  In either case, if invalid, the
        # user is re-prompted to enter a new destination.
        destdir = None
Paul McCarthy's avatar
Paul McCarthy committed
304
305
        if self.args.dest is not None: response = self.args.dest
        else:                          response = None
306
307
308
309

        while destdir is None:

            if response is None:
Paul McCarthy's avatar
Paul McCarthy committed
310
                printmsg('\nWhere do you want to install FSL?',
311
                         IMPORTANT, EMPHASIS)
312
313
314
315
                printmsg('Press enter to install to the default location '
                         '[{}]\n'.format(DEFAULT_INSTALLATION_DIRECTORY), INFO)
                response = prompt('FSL installation directory [{}]:'.format(
                    DEFAULT_INSTALLATION_DIRECTORY), QUESTION, EMPHASIS)
316
                response = response.rstrip(op.sep)
317

318
319
                if response == '':
                    response = DEFAULT_INSTALLATION_DIRECTORY
320

321
            response  = op.expanduser(op.expandvars(response))
322
323
324
325
326
327
328
329
330
331
            response  = op.abspath(response)
            parentdir = op.dirname(response)
            if op.exists(parentdir):
                destdir = response
            else:
                printmsg('Destination directory {} does not '
                         'exist!'.format(parentdir), ERROR)
                response = None

        self.__destdir = destdir
Paul McCarthy's avatar
Paul McCarthy committed
332
333
334
335
336
        return self.__destdir


    @property
    def need_admin(self):
337
338
339
        """Returns True if administrator privileges will be needed to install
        FSL.
        """
Paul McCarthy's avatar
Paul McCarthy committed
340
341
        if self.__need_admin is not None:
            return self.__need_admin
342
343
344
        parentdir = op.dirname(self.destdir)
        self.__need_admin = Context.check_need_admin(parentdir)
        return self.__need_admin
Paul McCarthy's avatar
Paul McCarthy committed
345
346
347
348


    @property
    def admin_password(self):
349
350
        """Returns the user's administrator password, prompting them if needed.
        """
Paul McCarthy's avatar
Paul McCarthy committed
351
352
353
354
355
356
357
358
359
        if self.__admin_password is not None:
            return self.__admin_password
        if self.__need_admin == False:
            return None
        self.__admin_password = Context.get_admin_password()


    @property
    def manifest(self):
360
        """Returns the FSL installer manifest as a dictionary. """
Paul McCarthy's avatar
Paul McCarthy committed
361
        if self.__manifest is None:
362
363
            self.__manifest = Context.download_manifest(self.args.manifest,
                                                        self.args.workdir)
Paul McCarthy's avatar
Paul McCarthy committed
364
365
366
367
368
369
370
371
372
373
374
375
376
377
        return self.__manifest


    @staticmethod
    def identify_platform():
        """Figures out what platform we are running on. Returns a platform
        identifier string - one of:

          - "linux-64" (Linux, x86_64)
          - "macos-64" (macOS, x86_64)
        """

        platforms = {
            ('linux',  'x86_64') : 'linux-64',
378
            ('darwin', 'x86_64') : 'macos-64',
379
380
381

            # M1 builds (and possbily ARM for Linux)
            # will be added in the future
Paul McCarthy's avatar
Paul McCarthy committed
382
383
384
385
386
387
388
389
            ('darwin', 'arm64')  : 'macos-64',
        }

        system = platform.system().lower()
        cpu    = platform.machine()
        key    = (system, cpu)

        if key not in platforms:
Paul McCarthy's avatar
Paul McCarthy committed
390
391
392
            supported = ', '.join(['[{}, {}]' for s, c in platforms])
            raise Exception('This platform [{}, {}] is unrecognised or '
                            'unsupported! Supported platforms: {}'.format(
393
                                system, cpu, supported))
Paul McCarthy's avatar
Paul McCarthy committed
394
395
396
397
398
399
400

        return platforms[key]


    @staticmethod
    def identify_cuda():
        """Identifies the CUDA version supported on the platform. Returns a
Paul McCarthy's avatar
Paul McCarthy committed
401
402
        float representing the X.Y CUDA version, or None if CUDA is not
        available on the platform.
Paul McCarthy's avatar
Paul McCarthy committed
403
404
        """

Paul McCarthy's avatar
Paul McCarthy committed
405
406
        # see below - no_cuda is set to prevent unnecessary
        # attempts to call nvidia-smi more than once
407
408
409
        if getattr(Context.identify_cuda, 'no_cuda', False):
            return None

Paul McCarthy's avatar
Paul McCarthy committed
410
        try:
411
            output = Process.check_output('nvidia-smi')
412
        except Exception:
413
            Context.identify_cuda.no_cuda = True
Paul McCarthy's avatar
Paul McCarthy committed
414
415
            return None

416
417
418
419
420
421
422
423
        pat   = r'CUDA Version: (\S+)'
        lines = output.split('\n')
        for line in lines:
            match = re.search(pat, line)
            if match:
                cudaver = match.group(1)
                break
        else:
Paul McCarthy's avatar
Paul McCarthy committed
424
425
426
            # message for debugging - the output
            # will be present in the logfile
            log.debug('Could not parse nvidia-smi output')
427
            Context.identify_cuda.no_cuda = True
428
            return None
Paul McCarthy's avatar
Paul McCarthy committed
429

Paul McCarthy's avatar
Paul McCarthy committed
430
        return float(cudaver)
Paul McCarthy's avatar
Paul McCarthy committed
431
432
433
434
435
436
437


    @staticmethod
    def check_need_admin(dirname):
        """Returns True if dirname needs administrator privileges to write to,
        False otherwise.
        """
438
439
        # os.supports_effective_ids added in
        # python 3.3, so can't be used here
Paul McCarthy's avatar
Paul McCarthy committed
440
441
442
443
444
445
446
447
        return not os.access(dirname, os.W_OK | os.X_OK)


    @staticmethod
    def get_admin_password():
        """Prompt the user for their administrator password."""

        def validate_admin_password(password):
448
            proc = Process.sudo_popen(['true'], password)
449
450
            proc.communicate()
            return proc.returncode == 0
Paul McCarthy's avatar
Paul McCarthy committed
451

Paul McCarthy's avatar
Paul McCarthy committed
452
453
454
455
456
457
458
        for attempt in range(3):
            if attempt == 0:
                msg = 'Your administrator password is needed to ' \
                      'install FSL: '
            else:
                msg = 'Your administrator password is needed to ' \
                      'install FSL [attempt {} of 3]:'.format(attempt + 1)
459
            printmsg(msg, IMPORTANT, end='')
Paul McCarthy's avatar
Paul McCarthy committed
460
461
462
            password = getpass.getpass('')
            valid    = validate_admin_password(password)

Paul McCarthy's avatar
Paul McCarthy committed
463
            if valid:
464
                printmsg('Password accepted', INFO)
Paul McCarthy's avatar
Paul McCarthy committed
465
466
467
                break
            else:
                printmsg('Incorrect password', WARNING)
Paul McCarthy's avatar
Paul McCarthy committed
468
469

        if not valid:
Paul McCarthy's avatar
Paul McCarthy committed
470
            raise Exception('Incorrect password')
Paul McCarthy's avatar
Paul McCarthy committed
471
472
473
474
475

        return password


    @staticmethod
476
    def download_manifest(url, workdir=None):
Paul McCarthy's avatar
Paul McCarthy committed
477
478
479
480
        """Downloads the installer manifest file, which contains information
        about available FSL vesrions, and the most recent version number of the
        installer (this script).

481
482
483
        The manifest file is a JSON file. Lines beginning
        with a double-forward-slash are ignored. See test/data/manifes.json
        for an example.
Paul McCarthy's avatar
Paul McCarthy committed
484
485
486

        This function modifies the manifest structure by adding a 'version'
        attribute to all FSL build entries.
Paul McCarthy's avatar
Paul McCarthy committed
487
488
489
490
        """

        log.debug('Downloading FSL installer manifest from %s', url)

491
        with tempdir(workdir):
Paul McCarthy's avatar
Paul McCarthy committed
492
493
            download_file(url, 'manifest.json')
            with open('manifest.json') as f:
494
495
496
497
                lines = f.readlines()

        # Drop comments
        lines = [l for l in lines if not l.lstrip().startswith('//')]
Paul McCarthy's avatar
Paul McCarthy committed
498

Paul McCarthy's avatar
Paul McCarthy committed
499
500
501
502
503
504
505
506
507
508
        manifest = json.loads('\n'.join(lines))

        # Add "version" to every build
        for version, builds in manifest['versions'].items():
            if version == 'latest':
                continue
            for build in builds:
                build['version'] = version

        return manifest
Paul McCarthy's avatar
Paul McCarthy committed
509
510


511
512
513
514
515
516
517
518
519
520
521
522
# List of modifiers which can be used to change how
# a message is printed by the printmsg function.
INFO      = 1
IMPORTANT = 2
QUESTION  = 3
PROMPT    = 4
WARNING   = 5
ERROR     = 6
EMPHASIS  = 7
UNDERLINE = 8
RESET     = 9
ANSICODES = {
523
524
525
526
527
528
529
530
531
    INFO      : '\033[37m',         # Light grey
    IMPORTANT : '\033[92m',         # Green
    QUESTION  : '\033[36m\033[4m',  # Blue+underline
    PROMPT    : '\033[36m\033[1m',  # Bright blue+bold
    WARNING   : '\033[93m',         # Yellow
    ERROR     : '\033[91m',         # Red
    EMPHASIS  : '\033[1m',          # White+bold
    UNDERLINE : '\033[4m',          # Underline
    RESET     : '\033[0m',          # Used internally
532
533
534
}


535
def printmsg(msg='', *msgtypes, **kwargs):
536
537
    """Prints msg according to the ANSI codes provided in msgtypes.
    All other keyword arguments are passed through to the print function.
Paul McCarthy's avatar
Paul McCarthy committed
538
539
540
541
542

    :arg msgtypes: Message types to control formatting
    :arg log:      If True (default), the message is logged.

    All other keyword arguments are passed to the built-in print function.
543
    """
Paul McCarthy's avatar
Paul McCarthy committed
544
    logmsg   = kwargs.pop('log', msg != '')
545
546
    msgcodes = [ANSICODES[t] for t in msgtypes]
    msgcodes = ''.join(msgcodes)
Paul McCarthy's avatar
Paul McCarthy committed
547
548
    if logmsg:
        log.debug(msg)
549
    print('{}{}{}'.format(msgcodes, msg, ANSICODES[RESET]), **kwargs)
550
    sys.stdout.flush()
551
552
553
554
555
556


def prompt(prompt, *msgtypes, **kwargs):
    """Prompts the user for some input. msgtypes and kwargs are passed
    through to the printmsg function.
    """
Paul McCarthy's avatar
Paul McCarthy committed
557
    printmsg(prompt, *msgtypes, end='', log=False, **kwargs)
558

Paul McCarthy's avatar
Paul McCarthy committed
559
560
561
562
563
564
    if PY2: response = raw_input(' ').strip()
    else:   response = input(    ' ').strip()

    log.debug('%s: %s', prompt, response)

    return response
565
566


567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
class Progress(object):
    """Simple progress reporter. Displays one of the following:

       - If both a value and total are provided, a progress bar is shown
       - If only a value is provided, a cumulative count is shown
       - If nothing is provided, a spinner is shown.

    Use as a context manager, and call the update method to report progress,
    e,g:

        with Progress('%') as p:
            for i in range(100):
                p.update(i + 1, 100)
    """

    def __init__(self,
                 label='',
                 transform=None,
                 fmt='{:.1f}',
                 total=None,
                 width=None):
        """Create a Progress reporter.

        :arg label:     Units (e.g. "MB", "%",)

        :arg transform: Function to transform values (see e.g.
                        Progress.bytes_to_mb)

        :arg fmt:       Template string used to format value / total.

        :arg total:     Maximum value - overrides the total value passed to
                        the update method.

        :arg width:     Maximum width, if a progress bar is displayed. Default
                        is to automatically infer the terminal width (see
                        Progress.get_terminal_width).
        """

        if transform is None:
            transform = Progress.default_transform

        self.width     = width
        self.fmt       = fmt.format
        self.total     = total
        self.label     = label
        self.transform = transform

        # used by the spin function
        self.__last_spin = None

    @staticmethod
    def default_transform(val, total):
        return val, total

    @staticmethod
    def bytes_to_mb(val, total):
        if val   is not None: val   = val   / 1048576
        if total is not None: total = total / 1048576
        return val, total

    @staticmethod
    def percent(val, total):
        if val is None or total is None:
            return val, total
        return 100 * (val / total), 100

    def __enter__(self):
        return self

    def __exit__(self, *args, **kwargs):
Paul McCarthy's avatar
Paul McCarthy committed
637
        printmsg(log=False)
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663

    def update(self, value=None, total=None):

        if total is None:
            total = self.total

        value, total = self.transform(value, total)

        if value is None and total is None:
            self.spin()
        elif value is not None and total is None:
            self.count(value)
        elif value is not None and total is not None:
            self.progress(value, total)

    def spin(self):

        symbols = ['|', '/', '-',  '\\']

        if self.__last_spin is not None: last = self.__last_spin
        else:                            last = symbols[-1]

        idx  = symbols.index(last)
        idx  = (idx + 1) % len(symbols)
        this = symbols[idx]

Paul McCarthy's avatar
Paul McCarthy committed
664
        printmsg(this, end='\r', log=False)
665
666
667
        self.__last_spin = this

    def count(self, value):
668

669
        value = self.fmt(value)
670
671
672
673

        if self.label is None: line = '{} ...'.format(value)
        else:                  line = '{}{} ...'.format(value, self.label)

Paul McCarthy's avatar
Paul McCarthy committed
674
        printmsg(line, end='\r', log=False)
675
676
677

    def progress(self, value, total):

678
        value = min(value, total)
679
680
681
682
683
684

        # arbitrary fallback of 50 columns if
        # terminal width cannot be determined
        if self.width is None: width = Progress.get_terminal_width(50)
        else:                  width = self.width

685
686
687
688
689
690
        fvalue = self.fmt(value)
        ftotal = self.fmt(total)
        suffix = '{} / {} {}'.format(fvalue, ftotal, self.label).rstrip()

        # +5: - square brackets around bar
        #     - space between bar and tally
691
        #     - space+spin at the end
692
        width     = width - (len(suffix) + 5)
693
694
        completed = int(round(width * (value  / total)))
        remaining = width - completed
695
696
697
        progress  = '[{}{}] {}'.format('#' * completed,
                                       ' ' * remaining,
                                       suffix)
698

Paul McCarthy's avatar
Paul McCarthy committed
699
700
        printmsg(progress, end='', log=False)
        printmsg(' ', end='', log=False)
701
        self.spin()
Paul McCarthy's avatar
Paul McCarthy committed
702
        printmsg(end='\r', log=False)
703
704
705
706
707
708
709


    @staticmethod
    def get_terminal_width(fallback=None):
        """Return the number of columns in the current terminal, or fallback
        if it cannot be determined.
        """
710
        # os.get_terminal_size added in python
711
        # 3.3, so we try it but fall back to tput
712
        try:
713
            return os.get_terminal_size()[0]
714
715
        except Exception:
            pass
716

717
        try:
Paul McCarthy's avatar
Paul McCarthy committed
718
            result = sp.check_output(('tput', 'cols'))
719
            return int(result.strip())
720
        except Exception:
721
722
723
            return fallback


724
@contextlib.contextmanager
725
726
727
728
729
730
def tempdir(override_dir=None):
    """Returns a context manager which creates, changes into, and returns a
    temporary directory, and then deletes it on exit.

    If override_dir is not None, instead of creating and changing into a
    temporary directory, this function just changes into override_dir.
731
732
    """

733
734
735
    if override_dir is None: tmpdir = tempfile.mkdtemp()
    else:                    tmpdir = override_dir

736
737
738
    prevdir = os.getcwd()

    try:
739
740
        os.chdir(tmpdir)
        yield tmpdir
741
742
743

    finally:
        os.chdir(prevdir)
744
745
        if override_dir is None:
            shutil.rmtree(tmpdir)
746
747


748
def sha256(filename, check_against=None, blocksize=1048576):
Paul McCarthy's avatar
Paul McCarthy committed
749
750
751
    """Calculate the SHA256 checksum of the given file. If check_against
    is provided, it is compared against the calculated checksum, and an
    error is raised if they are not the same.
752
753
    """

Paul McCarthy's avatar
Paul McCarthy committed
754
    hashobj = hashlib.sha256()
755

Paul McCarthy's avatar
Paul McCarthy committed
756
757
758
759
760
761
    with open(filename, 'rb') as f:
        while True:
            block = f.read(blocksize)
            if len(block) == 0:
                break
            hashobj.update(block)
762

Paul McCarthy's avatar
Paul McCarthy committed
763
764
    checksum = hashobj.hexdigest()

765
    if check_against is not None:
Paul McCarthy's avatar
Paul McCarthy committed
766
        if checksum != check_against:
Paul McCarthy's avatar
Paul McCarthy committed
767
768
            raise Exception('File {} does not match expected checksum '
                            '({})'.format(filename, check_against))
769
770

    return checksum
771
772


773
def download_file(url, destination, progress=None, blocksize=131072):
Paul McCarthy's avatar
Paul McCarthy committed
774
775
    """Download a file from url, saving it to destination. """

776
777
    def default_progress(downloaded, total):
        pass
Paul McCarthy's avatar
Paul McCarthy committed
778

779
780
781
782
    if progress is None:
        progress = default_progress

    log.debug('Downloading %s ...', url)
Paul McCarthy's avatar
Paul McCarthy committed
783

784
785
786
787
    # Path to local file
    if op.exists(url):
        url = 'file:' + urlrequest.pathname2url(op.abspath(url))

788
    req = None
Paul McCarthy's avatar
Paul McCarthy committed
789
    try:
790
791
792
793
        # py2: urlopen result cannot be
        # used as a context manager
        req = urlrequest.urlopen(url)
        with open(destination, 'wb') as outf:
Paul McCarthy's avatar
Paul McCarthy committed
794
795
796
797
798
799

            try:             total = int(req.headers['content-length'])
            except KeyError: total = None

            downloaded = 0

Paul McCarthy's avatar
Paul McCarthy committed
800
            progress(downloaded, total)
Paul McCarthy's avatar
Paul McCarthy committed
801
802
803
804
805
806
            while True:
                block = req.read(blocksize)
                if len(block) == 0:
                    break
                downloaded += len(block)
                outf.write(block)
807
                progress(downloaded, total)
Paul McCarthy's avatar
Paul McCarthy committed
808

809
    finally:
810
811
        if req:
            req.close()
Paul McCarthy's avatar
Paul McCarthy committed
812
813


814
class Process(object):
815
    """Container for a subprocess.Popen object, allowing non-blocking
816
817
818
819
820
821
822
    line-based access to its standard output and error streams via separate
    queues, while logging all outputs.

    Don't create a Process directly - use one of the following static methods:
     - Process.check_output
     - Process.check_call
     - Process.monitor_progress
823
824
    """

825

826
    def __init__(self, cmd, admin=False, ctx=None, log_output=True, **kwargs):
827
828
        """Run the specified command. Starts threads to capture stdout and
        stderr.
829

830
831
832
833
834
835
836
        :arg cmd:        Command to run - passed directly to subprocess.Popen
        :arg admin:      Run the command with administrative privileges
        :arg ctx:        The installer Context. Only used for admin password -
                         can be None if admin is False.
        :arg log_output: If True, the command and all of its stdout/stderr are
                         logged.
        :arg kwargs:     Passed to subprocess.Popen
837
        """
838

839
840
841
842
843
844
845
846
847
848
        self.ctx        = ctx
        self.cmd        = cmd
        self.admin      = admin
        self.log_output = log_output
        self.stdoutq    = queue.Queue()
        self.stderrq    = queue.Queue()

        if log_output:
            log.debug('Running %s [as admin: %s]', cmd, admin)

849
        self.popen = Process.popen(self.cmd, self.admin, self.ctx, **kwargs)
850

851
        # threads for consuming stdout/stderr
852
853
        self.stdout_thread = threading.Thread(
            target=Process.forward_stream,
854
            args=(self.popen.stdout, self.stdoutq, cmd, 'stdout', log_output))
855
856
        self.stderr_thread = threading.Thread(
            target=Process.forward_stream,
857
            args=(self.popen.stderr, self.stderrq, cmd, 'stderr', log_output))
858

859
860
861
862
        self.stdout_thread.daemon = True
        self.stderr_thread.daemon = True
        self.stdout_thread.start()
        self.stderr_thread.start()
863
864


865
866
867
868
869
870
871
872
873
    def wait(self):
        """Waits for the process to terminate, then waits for the stdout
        and stderr consumer threads to finish.
        """
        self.popen.wait()
        self.stdout_thread.join()
        self.stderr_thread.join()


874
875
876
877
878
879
880
881
882
883
884
    @property
    def returncode(self):
        """Process return code. Returns None until the process has terminated,
        and the stdout/stderr consumer threads have finished.
        """
        if self.popen.returncode is None: return None
        if self.stdout_thread.is_alive(): return None
        if self.stderr_thread.is_alive(): return None
        return self.popen.returncode


885
    @staticmethod
886
    def check_output(cmd, *args, **kwargs):
887
888
889
        """Behaves like subprocess.check_output. Runs the given command, then
        waits until it finishes, and return its standard output. An error
        is raised if the process returns a non-zero exit code.
890

891
        :arg cmd: The command to run, as a string
892
893
        """

894
        proc = Process(cmd, *args, **kwargs)
895
        proc.wait()
896

897
        if proc.returncode != 0:
Paul McCarthy's avatar
Paul McCarthy committed
898
            raise RuntimeError('This command returned an error: ' + cmd)
899

900
901
902
903
904
905
906
907
        stdout = ''
        while True:
            try:
                stdout += proc.stdoutq.get_nowait()
            except queue.Empty:
                break

        return stdout
908
909
910


    @staticmethod
911
    def check_call(cmd, *args, **kwargs):
912
913
914
915
        """Behaves like subprocess.check_call. Runs the given command, then
        waits until it finishes. An error is raised if the process returns a
        non-zero exit code.

916
        :arg cmd: The command to run, as a string
917
        """
918
        proc = Process(cmd, *args, **kwargs)
919
        proc.wait()
920
        if proc.returncode != 0:
Paul McCarthy's avatar
Paul McCarthy committed
921
            raise RuntimeError('This command returned an error: ' + cmd)
922
923
924


    @staticmethod
925
    def monitor_progress(cmd, total=None, *args, **kwargs):
926
        """Runs the given command(s), and shows a progress bar under the
927
        assumption that cmd will produce "total" number of lines of output.
928
929
930
931

        :arg cmd:   The commmand to run as a string, or a sequence of
                    multiple commands.
        :arg total: Total number of lines of standard output to expect.
932
933
934
935
        """
        if total is None: label = None
        else:             label = '%'

936
937
938
        if isinstance(cmd, str): cmds = [cmd]
        else:                    cmds =  cmd

939
940
941
942
943
        with Progress(label=label,
                      fmt='{:.0f}',
                      transform=Progress.percent) as prog:


944
            for cmd in cmds:
945

946
947
                proc   = Process(cmd, *args, **kwargs)
                nlines = 0 if total else None
948
949
950

                prog.update(nlines, total)

951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
                while proc.returncode is None:
                    try:
                        line    = proc.stdoutq.get(timeout=0.5)
                        nlines  = (nlines + 1) if total else None

                    except queue.Empty:
                        pass

                    prog.update(nlines, total)
                    proc.popen.poll()

                # force progress bar to 100% when finished
                if proc.returncode == 0:
                    prog.update(total, total)
                else:
                    raise RuntimeError('This command returned '
                                       'an error: ' + cmd)
968

969
970

    @staticmethod
971
    def forward_stream(stream, queue, cmd, streamname, log_output):
972
973
974
        """Reads lines from stream and pushes them onto queue until popen
        is finished. Logs every line.

975
        :arg stream:     stream to forward
976
977
978
        :arg queue:      queue.Queue to push lines onto
        :arg cmd:        string - the command that is running
        :arg streamname: string - 'stdout' or 'stderr'
979
        :arg log_output: If True, log all stdout/stderr.
980
981
        """

982
        while True:
983
984
985
986
987
            line = stream.readline().decode('utf-8')
            if line == '':
                break
            else:
                queue.put(line)
988
                if log_output:
989
                    log.debug(' [%s]: %s', streamname, line.rstrip())
990
991
992


    @staticmethod
993
    def popen(cmd, admin=False, ctx=None, **kwargs):
994
995
        """Runs the given command via subprocess.Popen, as administrator if
        requested.
996

997
998
999
        :arg cmd:    The command to run, as a string

        :arg admin:  Whether to run with administrative privileges
1000

1001
1002
        :arg ctx:    The installer Context object. Only required if admin is
                     True.
1003

1004
1005
        :arg kwargs: Passed to subprocess.Popen. stdin, stdout, and stderr
                     will be silently clobbered
1006

1007
        :returns:    The subprocess.Popen object.
1008
1009
1010
1011
1012
1013
1014
        """

        admin = admin and os.getuid() != 0

        if admin: password = ctx.password
        else:     password = None

1015
1016
1017
1018
        cmd              = shlex.split(cmd)
        kwargs['stdin']  = sp.PIPE
        kwargs['stdout'] = sp.PIPE
        kwargs['stderr'] = sp.PIPE
1019
1020

        if admin: proc = Process.sudo_popen(cmd, password, **kwargs)
1021
        else:     proc = sp.Popen(          cmd,           **kwargs)
1022
1023
1024
1025
1026

        return proc


    @staticmethod
1027
1028
1029
    def sudo_popen(cmd, password, **kwargs):
        """Runs "sudo cmd" using subprocess.Popen. Used by Process.popen.
        Assumes that kwargs contains stdin=sp.PIPE
1030
        """
1031

1032
1033
1034
1035
        cmd  = ['sudo', '-S', '-k'] + cmd
        proc = sp.Popen(cmd, **kwargs)
        proc.stdin.write('{}\n'.format(password))
        return proc
1036
1037


1038
def list_available_versions(manifest):
1039
1040
    """Lists available FSL versions. """
    printmsg('Available FSL versions:', EMPHASIS)
1041
    for version in manifest['versions']:
1042
1043
1044
        if version == 'latest':
            continue
        printmsg(version, IMPORTANT, EMPHASIS)
1045
        for build in manifest['versions'][version]:
1046
1047
1048
1049
1050
1051
1052
1053
            if build.get('cuda', '').strip() != '':
                template = '  {platform} [CUDA {cuda}]'
            else:
                template = '  {platform}'
            printmsg(template.format(**build), EMPHASIS, end=' ')
            printmsg(build['environment'], INFO)


1054
1055
1056
1057
def download_fsl_environment(ctx):
    """Downloads the environment specification file for the selected FSL
    version.

1058
1059
1060
    If the (hidden) --environment option is provided, the specified file
    is used instead.

1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
    Internal/development FSL versions may source packages from the internal
    FSL conda channel, which requires a username+password to authenticate.

    These are referred to in the environment file as ${FSLCONDA_USERNAME}
    and ${FSLCONDA_PASSWORD}.

    If the user has not provided a username+password on the command-line, they
    are prompted for them.
    """

1071
1072
1073
1074
1075
1076
1077
1078
    if ctx.args.environment is None:
        build    = ctx.build
        url      = build['environment']
        checksum = build['sha256']
    else:
        build    = {}
        url      = ctx.args.environment
        checksum = None
1079
1080
1081
1082
1083
1084

    printmsg('Downloading FSL environment specification '
             'from {}...'.format(url))
    fname = url.split('/')[-1]
    download_file(url, fname)
    ctx.environment_file = op.abspath(fname)
1085
1086
    if (checksum is not None) and (not ctx.args.no_checksum):
        sha256(fname, checksum)
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114

    # Environment files for internal/dev FSL versions
    # will list the internal FSL conda channel with
    # ${FSLCONDA_USERNAME} and ${FSLCONDA_PASSWORD}
    # as placeholders for the username/password.
    with open(fname, 'rt') as f:
        need_auth = '${FSLCONDA_USERNAME}' in f.read()

    # We need a username/password to access the internal
    # FSL conda channel. Prompt the user if they haven't
    # provided credentials.
    if need_auth and (ctx.args.username is None):
        printmsg('A username and password are required to install '
                 'this version of FSL.', WARNING, EMPHASIS)
        ctx.args.username = prompt('Username:').strip()
        ctx.args.password = getpass.getpass('Password: ').strip()

    # Conda expands environment variables within a
    # .condarc file, but *not* within an environment.yml
    # file. So to authenticate to our internal channel
    # without storing credentials anywhere in plain text,
    # we *move* the channel list from the environment.yml
    # file into $FSLDIR/.condarc.
    #
    # Here we extract the channels from the environment
    # file, and save them to ctx.environment_channels.
    # The install_miniconda function will then add the
    # channels to $FSLDIR/.condarc.
1115
1116
1117
1118
1119
1120
    #
    # The fsl-base version is installed before any other
    # packages, as other FSL packages require it to be
    # present in order to install successfully. So we
    # also extract the fsl-base vesrion number from
    # the environment file, and store it in the context.
1121
    channels = []
Paul McCarthy's avatar
Paul McCarthy committed
1122
    basever  = None
1123
    copy     = '.' + op.basename(ctx.environment_file)
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
    shutil.move(ctx.environment_file, copy)
    with open(copy,                 'rt') as inf, \
         open(ctx.environment_file, 'wt') as outf:

        in_channels_section = False

        for line in inf:

            # start of channels list
            if line.strip() == 'channels:':
                in_channels_section = True
                continue

            if in_channels_section:
                # end of channels list
                if not line.strip().startswith('-'):
                    in_channels_section = False
                else:
                    channels.append(line.split()[-1])
                    continue

1145
1146
1147
1148
1149
            # save fsl-base version, as
            # we install it separately
            if line.strip().startswith('- fsl-base'):
                basever = line.split()[2]

1150
1151
1152
            outf.write(line)

    ctx.environment_channels = channels
1153
    ctx.fsl_base_version     = basever
1154
1155


Paul McCarthy's avatar
Paul McCarthy committed
1156
1157
1158
def download_miniconda(ctx):
    """Downloads the miniconda/miniforge installer and saves it as
    "miniconda.sh".
1159
1160

    This function assumes that it is run within a temporary/scratch directory.
1161
1162
    """

1163
    metadata = ctx.manifest['miniconda'][ctx.platform]
1164
1165
1166
    url      = metadata['url']
    checksum = metadata['sha256']
    output   = metadata.get('output', '').strip()
1167

1168
1169
    if output == '': output = None
    else:            output = int(output)
Paul McCarthy's avatar
Paul McCarthy committed
1170

1171
1172
    # Download
    printmsg('Downloading miniconda from {}...'.format(url))
1173
    with Progress('MB', transform=Progress.bytes_to_mb) as prog:
1174
        download_file(url, 'miniconda.sh', prog.update)
1175
    if not ctx.args.no_checksum:
Paul McCarthy's avatar
Paul McCarthy committed
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
        sha256('miniconda.sh', checksum)


def install_miniconda(ctx):
    """Downloads the miniconda/miniforge installer, and installs it to the
    destination directory.

    This function assumes that it is run within a temporary/scratch directory.
    """

    metadata = ctx.manifest['miniconda'][ctx.platform]
    url      = metadata['url']
    checksum = metadata['sha256']
    output   = metadata.get('output', '').strip()

    if output == '': output = None
    else:            output = int(output)
1193
1194

    # Install
1195
    printmsg('Installing miniconda at {}...'.format(ctx.destdir))
Paul McCarthy's avatar
Paul McCarthy committed
1196
    cmd = 'sh miniconda.sh -b -p {}'.format(ctx.destdir)
1197
    Process.monitor_progress(cmd, output, ctx.need_admin, ctx)
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213

    # Create .condarc config file
    condarc = tw.dedent("""
    # Try and make package downloads more robust
    remote_read_timeout_secs:    240
    remote_connect_timeout_secs: 20
    remote_max_retries:          10
    remote_backoff_factor:       5
    safety_checks:               warn

    # Channel priority is important. In older versions
    # of FSL we placed the FSL conda channel at the
    # bottom (lowest priority) for legacy reasons (to
    # ensure that conda-forge versions of e.g. VTK were
    # preferred over legacy FSL conda versions).
    #
1214
1215
1216
1217
1218
1219
    # Use final/top/bottom marks to prevent the channel
    # priority order being modified by user ~/.condarc
    # configuration files.
    #
    # https://conda.io/projects/conda/en/latest/user-guide/configuration/use-condarc.html
    # https://www.anaconda.com/blog/conda-configuration-engine-power-users
1220
    # https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-channels.html
1221
    channel_priority: strict #!final
1222
    """)
1223
    channels      = list(ctx.environment_channels)
1224
1225
1226
    if len(channels) > 0:
        channels[0]  += ' #!top'
        channels[-1] += ' #!bottom'
1227
1228
    condarc      += '\nchannels: #!final\n'
    for channel in channels:
1229
        condarc += ' - {}\n'.format(channel)
1230
1231
1232
1233

    with open('.condarc', 'wt') as f:
        f.write(condarc)

1234
1235
    Process.check_call('cp .condarc {}'.format(ctx.destdir),
                       ctx.need_admin, ctx)