#!/usr/bin/python
"""olpc-update installs a new version of the XO's base filesystem."""
# Copyright (C) 2007-8 One Laptop Per Child Association, Inc.
# Licensed under the terms of the GNU GPL v2 or later; see COPYING for details.
# Written by C. Scott Ananian <cscott@laptop.org>
from __future__ import with_statement

from tempfile import mkdtemp
from subprocess import call, check_call, CalledProcessError
from bitfrost.contents.utils import UnifiedContents
from bitfrost.update.setup import shorten_hash
from bitfrost.update import perform_update, inhibit_suspend, VERSION_INFO
import bitfrost.util.json as json
import os, os.path, sys
from shutil import copyfile, rmtree

def main():
    from optparse import OptionParser
    parser = OptionParser(usage="""
 %prog [options] --hints hints-file
 %prog [options] --latest stream-name
 %prog [options] [-rf] build-number
 %prog [options] [-rf] --usb
 %prog --version
 %prog --help

For example:
 %prog 630
 %prog joyride-258
 %prog --latest joyride

For convenience, '%prog joyride-latest' is treated as an alias for
'%prog --latest joyride'.""")
    parser.add_option('-f','--full',action='store_true',dest='full',
                      help='skip incremental update attempt.')
    parser.add_option('--force', action='store_true', dest='force',
                      help='force update to an unsigned build.')
    parser.add_option('-r','--reboot',action='store_true',dest='reboot',
                      help='reboot after a successful update.')
    parser.add_option('-p','--poweroff',action='store_true',dest='poweroff',
                      help='poweroff after a successful update.')
    parser.add_option('--hints',dest='hints',default=None,metavar='FILE',
                      help='name of json-encoded hints dictionary identifying '+
                      'the desired new version.')
    parser.add_option('--latest',dest='stream',default=None,metavar='STREAM',
                      help='update to the version named by the given update '+
                      'stream.')
    parser.add_option('-u','--usb',action='store_true',dest='usb',
                      help='update to new build on inserted USB stick.')
    parser.add_option('-s','--server', dest='server', default='updates.laptop.org',
                      help='alternative update server')
    parser.add_option('-v',action='count',dest='verbosity',default=0,
                      help='display verbose progress information; repeat for more verbose output.')
    parser.add_option('-q','--quiet',action='store_true',dest='quiet',
                      default=False, help="don't output anything; "+
                      "use exit status to indicate success.")
    parser.add_option('--version',action='store_true',dest='version',
                      default=False,
                      help="display version and license information.")
    (options, args) = parser.parse_args()
    if options.version:
        print VERSION_INFO
    if os.getuid() != 0:
        parser.error('Must be run as root.')
    no_expected_args = (options.hints is not None or options.stream is not None
                        or options.usb)
    if len(args) == 0 and not no_expected_args:
        parser.error('Build number is required.')
    if len(args) > (0 if no_expected_args else 1):
        parser.error('Too many arguments.')
    # convenience shortcut: build number ending in '-latest' maps to
    # equivalent '--latest' invocation.
    if not no_expected_args and args[0].endswith('-latest'):
        options.stream = args[0].rsplit('-',1)[0]
    olpc_update(options, args)

@inhibit_suspend
def olpc_update(options, args):
    def report(lvl, msg):
        if options.quiet: return
        if options.verbosity >= lvl:
            print >>sys.stderr, msg

    # look up current hash.
    try:
        cur_hash = os.path.basename(os.readlink('/versions/running'))
    except OSError: # allow testing on non-XO
        report(0, "No current running version?  Continuing anyway...")
        cur_hash = 'xyzzy'

    # special case --latest, which is really shorthand for an invocation of
    # olpc-update-query
    if options.stream is not None:
        verbosity=['-v']*(options.verbosity+1) # query defaults to verbosity=-1
        retcode = call(['olpc-update-query', '-f', '--stream', options.stream]+
                       verbosity)
        sys.exit(retcode)

    # create an appropriate hints file.
    if options.hints is not None:
        new_hash, priority, hints = json.read(open(options.hints).read())
    elif options.usb:
        import glob
        # Find an osXYZ.toc / osXYZ.usb pair in /media/*
        candidates = glob.glob('/media/*/os*.usb')
        if len(candidates)==0:
            report(-1, "Can't find os*.usb file on mounted USB key.")
            sys.exit(1)
        elif len(candidates)>1:
            report(-1, "Multiple os*.usb files found: %s" %
                   (', '.join(candidates)))
            report(-1, "I'm confused. There should be only one.")
            sys.exit(1)
        usb_file = candidates[0]
        toc_file = usb_file[:-3]+'toc'
        if not os.path.exists(toc_file):
            report(-1, "Can't find %s.  Halting." % toc_file)
            sys.exit(1)
        # Get hash from the toc file.
        new_hash = UnifiedContents(toc_file).contents_hash()
        # Set priority.
        priority = 'urgent' if options.reboot or options.poweroff else 'normal'
        # Create a working directory $tmpdir
        tmpdir = mkdtemp()
        try:
            # mount .usb file in $tmpdir/root
            root = os.path.join(tmpdir, 'root')
            os.mkdir(root)
            check_call(['/bin/mount','-o','loop,ro',usb_file,root])
            try:
                # copy .toc file to $tmpdir/contents
                copyfile(toc_file, os.path.join(tmpdir, 'contents'))
                # make hints
                url = tmpdir
                hints = [('irsync_pristine', url),
                         ('irsync_dirty', url),
                         ('rsync', url),
                         ] if not options.full else [('rsync', url)]
                hints_file = os.path.join(tmpdir, 'hints')
                with open(hints_file, 'w') as f:
                    f.write(json.write([new_hash, priority, hints]))
                # recurse!
                options.usb = False
                options.hints = hints_file
                olpc_update(options, args)
                assert False # should never reach this point.
            finally:
                # unmount fs.
                check_call(['/bin/umount', root])
        finally:
            # cleanup tmpdir
            rmtree(tmpdir)
    else:
        build = args[0]
        # sanity check: do we need a signed build (trac #5307, #7265)
        if ('-' in build and not build.startswith('candidate-')) \
               and os.path.exists('/ofw/mfg-data/wp') and \
               not (os.path.exists('/security/develop.sig') or
                    os.path.exists('/ofw/mfg-data/dk')):
            report(-1, "WARNING: You seem to be attempting to download "+
                   "an unsigned development build,\nbut you don't have a "+
                   "developer key.  This will probably fail.")
        
        # download and hash the contents file to get hash.  not pretty.
        tmpdir = mkdtemp()
        new_contents = os.path.join(tmpdir, 'contents')
        try: # try to seed rsync with the existing contents file for speed.
            copyfile('/versions/contents/%s' % cur_hash, new_contents)
        except: pass # if it doesn't work, no worries.
        try:
            report(0, "Downloading contents of build %s." % build)
            url = 'rsync://%s/build-%s' % (options.server, build)
            check_call(['rsync', '-z', url+'/contents', new_contents])
            # hash the relevant bit of the contents file.
            c = UnifiedContents(new_contents)

            new_hash, priority, hints = \
              ( c.contents_hash(),
                'urgent' if options.reboot or options.poweroff else 'normal',
                [('irsync_pristine', url),
                 ('irsync_dirty', url),
                 ('rsync', url),
                 ] if not options.full else [('rsync', url)])
        except CalledProcessError, detail:
            report(-1, '')
            report(-1, "Could not download update contents file from:")
            report(-1, "  %s/contents" % url)
            if detail.returncode == 5:
                report(-1, "I don't think the requested build number exists.")
            elif detail.returncode == 10:
                report(-1, "I don't think you're connected to the internet.")
            else:
                report(-1, "Does this build exist?  Are you connected to the internet?")
            sys.exit(1)
        finally:
            rmtree(tmpdir)
    # maybe we're already up-to-date
    if shorten_hash(new_hash) == shorten_hash(cur_hash):
        if not options.quiet: print 'Already up-to-date.'
        sys.exit(0)
    # okay, now we've got the desired hash, priority, and update hints.
    # let's do the update!
    report(0, "Updating to version hash %s" % shorten_hash(new_hash))
    perform_update(new_hash, priority, hints,
                   options.verbosity if not options.quiet else -1,
                   force=options.force, poweroff=options.poweroff)
    assert False  # perform_update should never return
    
if __name__ == '__main__': main ()
