~bzoltan/phablet-tools/single_quotation_for_spaced_SSID

#! /usr/bin/python2.7
# -*- Mode: Python; coding: utf-8; indent-tabs-mode: nil; tab-width: 4 -*-
# Copyright (C) 2013 Canonical Ltd.
# Author: Chris Wayne <cwayne@ubuntu.com>

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 3 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import print_function

from phabletutils.device import AndroidBridge
from phabletutils.fileutils import create_temp_dir
from subprocess import CalledProcessError

import argparse
import os
import shutil
import subprocess
import sys
import urllib
import re
import pipes

LOCAL_SOURCE_FILE = '''deb file:/tmp/archive/ ./
'''
LOCAL_PIN_FILE = '''Package: *
Pin: origin ""
Pin-Priority: 1100
'''


class UbuntuDevice(AndroidBridge):

    def __init__(self, device=None):
        super(UbuntuDevice, self).__init__(device=device)


def _handle_timezone(adb, args):
    if not args.tz:
        try:
            tz_string = open('/etc/timezone', 'r').read().strip('\n')
            print(tz_string)
        except IOError:
            print("/etc/timezone not found.  Please pass --tz")
            sys.exit(1)
    else:
        tz_string = args.tz
    if not is_valid_tz(tz_string):
        print("Not a valid timezone")
        sys.exit(1)
    print("Changing timezone to %s" % tz_string)
    try:
        cmd = "gdbus call -y -d org.freedesktop.timedate1 -o " \
              "/org/freedesktop/timedate1 -m " \
              "org.freedesktop.timedate1.SetTimezone '%s' false" % tz_string
        print(cmd)
        adb.shell(cmd)
        print("Time zone successfully changed")
    except:
        print("Timezone has not been changed")
        sys.exit(1)


def is_valid_tz(timezone):
    tz_path = os.path.join('/usr/share/zoneinfo/' + timezone)
    return os.path.isfile(tz_path)


def is_valid_locale(adb, locale):
    phone_locales = adb.shell('ls /usr/share/locale/').split()
    return locale in phone_locales


def _handle_locale(adb, args):
    if not args.lc and not args.list:
        print("No locale to set, exiting")
        sys.exit(1)
    if args.list:
        print(adb.shell('ls /usr/share/locale'))
        sys.exit(0)
    elif not is_valid_locale(adb, args.lc):
        print("Not a valid locale, exiting")
        sys.exit(1)
    try:
        cmd = "dbus-send --system --print-reply --dest="\
              "org.freedesktop.Accounts /org/freedesktop/Accounts/User32011 "\
              "org.freedesktop.Accounts.User.SetLanguage"\
              " string:%s" % args.lc
        adb.shell(cmd)
        print("Locale set to %s, rebooting to take effect" % args.lc)
        adb.reboot()
    except:
        print("Locale was not changed")
        sys.exit(1)


def is_remote_root(adb):
    return adb.shell("id -ru").strip() == "0"


def _handle_autopilot(adb, args):
    if not is_remote_root(adb):
        dbusarg = 'false'
        if args.dbus_probe == 'enable':
            dbusarg = 'true'
        dbus_call = 'dbus-send --system --print-reply ' \
                    '--dest=com.canonical.PropertyService ' \
                    '/com/canonical/PropertyService ' \
                    'com.canonical.PropertyService.SetProperty ' \
                    'string:autopilot boolean:%s' % dbusarg
        adb.shell(dbus_call, False)
        # wait until it is done; it takes a while to start apparmor_parser, so
        # we need to wait a bit in advance
        adb.shell('sleep 5; while pidof apparmor_parser; do sleep 5; done',
                  False)
    else:
        if args.dbus_probe == 'enable':
            rfile = '/usr/share/autopilot-touch/apparmor/click.rules'
            adb.shell('aa-clickhook -f --include=%s' % rfile)
        else:
            adb.shell('aa-clickhook -f')


def sudo_shell(adb, sudocmd, cmd):
    output = adb.shell('%s %s' % (sudocmd, cmd), ignore_errors=False)
    print('%s' % output)


def _handle_writable_image(adb, args):
    fname = '/userdata/.writable_image'
    try:
        adb.shell('test -e %s' % fname, False)
    except CalledProcessError as e:
        if not is_remote_root(adb):
            dbus_call = 'dbus-send --system --print-reply ' \
                        '--dest=com.canonical.PropertyService ' \
                        '/com/canonical/PropertyService ' \
                        'com.canonical.PropertyService.SetProperty ' \
                        'string:writable boolean:true'
            adb.shell(dbus_call, False)
        else:
            adb.shell('touch %s' % fname, False)
        adb.reboot()

    adb.wait_for_device()
    adb.wait_for_network()

    try:
        sudocmd = ''
        tempfile = ''
        if not is_remote_root(adb) and args.remotepassword:
            tempfile = adb.shell('mktemp -t sudo_askpass.XXXX').rstrip()

            adb.shell("printf '%%s\\n' '#!/bin/sh' 'echo %s' >> %s " %
                      (args.remotepassword[0], tempfile), ignore_errors=False)
            adb.shell("chmod u+x %s" % tempfile)
            sudocmd = 'SUDO_ASKPASS=%s sudo -A' % tempfile

        if args.package_dir:
            if not is_remote_root(adb) and not args.remotepassword:
                print('need the remote password (-r) for package operations')
                sys.exit(1)
            tmp_archive = create_temp_dir()
            for package_dir in args.package_dir:
                # Copy all deb files under the given dir to a temp dir
                for root, dirs, files in os.walk(package_dir):
                    for deb_file in files:
                        if deb_file.endswith('.deb'):
                            shutil.copy(os.path.join(root, deb_file),
                                        tmp_archive)
            # Create the Packages file
            current_dir = os.getcwd()
            os.chdir(tmp_archive)
            subprocess.check_call(
                '/usr/bin/dpkg-scanpackages . /dev/null > Packages',
                shell=True)
            os.chdir(current_dir)
            # Create the apt source file for the device
            with open(os.path.join(tmp_archive,
                                   '00localrepo.list'), 'w') as sources_file:
                sources_file.write(LOCAL_SOURCE_FILE)
            # Create the apt pin file for the device
            with open(os.path.join(tmp_archive,
                                   'local-pin-1100'), 'w') as pin_file:
                pin_file.write(LOCAL_PIN_FILE)

            adb.push(tmp_archive, '/tmp/archive')
            sudo_shell(adb, sudocmd, 'mv /tmp/archive/00localrepo.list '
                       '/etc/apt/sources.list.d')
            sudo_shell(adb, sudocmd, 'mv /tmp/archive/local-pin-1100 '
                       '/etc/apt/preferences.d')

        if args.ppa:
            if not is_remote_root(adb) and not args.remotepassword:
                print('need the remote password (-r) for package operations')
                sys.exit(1)
            for ppa in args.ppa:
                sudo_shell(adb, sudocmd, 'add-apt-repository -y %s' % ppa)

        if args.package_dir or args.ppa:
            sudo_shell(adb, sudocmd, 'apt-get update -qq')

        if args.package:
            if not is_remote_root(adb) and not args.remotepassword:
                print('need the remote password (-r) for package operations')
                sys.exit(1)
            package = ' '.join(args.package)
            sudo_shell(adb, sudocmd,
                       'apt-get install -yq --force-yes %s' % package)

        if args.package_dir:
            # Remove the local archive
            sudo_shell(adb, sudocmd,
                       'rm /etc/apt/sources.list.d/00localrepo.list')
            sudo_shell(adb, sudocmd,
                       'rm /etc/apt/preferences.d/local-pin-1100')
            adb.shell('rm -rf /tmp/archive')
        adb.shell('rm -rf %s' % tempfile)

    except CalledProcessError as e:
        print(e)
        print(e.output)
        sys.exit(e.returncode)


def _handle_edges_intro(adb, args):
    fmt = 'dbus-send --system --print-reply --dest=org.freedesktop.Accounts ' \
          '/org/freedesktop/Accounts/User32011 %s ' \
          'string:com.canonical.unity.AccountsService string:demo-edges'
    getter = fmt % 'org.freedesktop.DBus.Properties.Get'

    enabled = adb.shell(getter).find('true') != -1
    if args.enable and enabled:
        print('already enabled')
    elif not args.enable and not enabled:
        print('already disabled')
    else:
        setter = 'true' if args.enable else 'false'
        dbus_call = 'dbus-send --system --print-reply ' \
                    '--dest=com.canonical.PropertyService ' \
                    '/com/canonical/PropertyService ' \
                    'com.canonical.PropertyService.SetProperty ' \
                    'string:edge boolean:%s' % setter
        adb.shell(dbus_call, False)


def _handle_welcome_wizard(adb, args):
    config_dir = '/home/phablet/.config/ubuntu-system-settings'
    ran_file = config_dir + '/wizard-has-run'

    enabled = adb.shell('ls %s' % ran_file).strip() != ran_file
    if args.enable and enabled:
        print('already enabled')
    elif not args.enable and not enabled:
        print('already disabled')
    elif args.enable:
        adb.shell('rm %s' % ran_file)
    else:
        adb.shell('sudo -u phablet mkdir -p %s' % config_dir)
        adb.shell('sudo -u phablet touch %s' % ran_file)


def get_wifi(adb):
    '''Return (ssid, password) of current wifi connection

    Raise ValueError if there is no current connection.
    '''
    out = adb.shell(
        'gdbus call --system -d org.freedesktop.NetworkManager '
        '-o /org/freedesktop/NetworkManager/ActiveConnection/0 '
        '-m org.freedesktop.DBus.Properties.Get '
        'org.freedesktop.NetworkManager.Connection.Active Connection')
    m = re.search("objectpath '(.*?)'>", out)
    if not m:
        raise ValueError('no current connection')
    settings_path = m.group(1)

    # get ESSID and secret type
    out = adb.shell(
        'gdbus call --system -d org.freedesktop.NetworkManager -o %s '
        '-m org.freedesktop.NetworkManager.Settings.Connection.GetSettings' %
        settings_path)
    m = re.search("'ssid': <\[byte (.*?)\]", out)
    assert m, 'failed to parse SSID out of %s' % out
    # convert byte array to string
    ssid = bytearray()
    for b in m.group(1).split(','):
        ssid.append(int(b, 16))
    m = re.search("'security': <'(.*?)'", out)
    if m:
        secret_type = m.group(1)
        out = adb.shell(
            'gdbus call --system -d org.freedesktop.NetworkManager -o %s '
            '-m org.freedesktop.NetworkManager.Settings.Connection.GetSecrets'
            ' %s' % (settings_path, secret_type))
        m = re.search("{'.*?': <'(.*?)'>}", out)
        assert m, 'failed to parse secret out of %s' % out
        secret = m.group(1)
    else:
        secret = ''

    return (ssid, secret)


def set_wifi(adb, ssid, password):
    '''Connect to given wifi'''

    out = adb.shell('nmcli device wifi connect %s password %s' % (
        pipes.quote(ssid), pipes.quote(password)))
    print(out)
    if 'Error' in out:
        sys.exit(1)


def _handle_network(adb, args):
    if args.read:
        try:
            (ssid, password) = get_wifi(adb)
            print('ssid=%s password=%s' %
                  (urllib.quote(ssid), urllib.quote(password)))
        except ValueError:
            sys.stderr.write('No current wifi connection\n')
            sys.exit(1)

    elif args.write:
        m = re.match('\s*ssid=([^\s]+)\s+password=([^\s]*)\s*$', args.write)
        if not m:
            sys.stderr.write('Invalid wifi specification format\n')
            sys.exit(1)
        set_wifi(adb, urllib.unquote(m.group(1)), urllib.unquote(m.group(2)))


def parse_arguments():
    parser = argparse.ArgumentParser(
        description='Set up different configuration options on a device.')
    parser.add_argument('-s', '--serial',
                        help='''Device serial. Use when more than
                                one device is connected.''')

    sub = parser.add_subparsers(title='Commands', metavar='')
    tz = sub.add_parser('timezone', help='Configure timezone for device.')
    tz.add_argument('--tz',
                    help='''Timezone to push to phone. Defaults to whatever
                            is in hosts /etc/timezone''')
    tz.set_defaults(func=_handle_timezone)

    lc = sub.add_parser('locale', help="Configure locale for device.")
    lc.add_argument('--lc', help='''Locale to push to phone.  See
                                    --list for valid locales''')
    lc.add_argument('--list', help='List available locales on device',
                    action='store_true')
    lc.set_defaults(func=_handle_locale)

    ap = sub.add_parser('autopilot', help='Configure autopilot for device.')
    ap.add_argument('--dbus-probe', choices=('enable', 'disable'),
                    required=True,
                    help='''Allows autopilot the ability to introspect DBUS
                         objects so that tests will work.''')
    ap.set_defaults(func=_handle_autopilot)
    dm = sub.add_parser('writable-image', help='''Enable read/write access to
                        device partitions  and optionally install extra
                        packages.''')
    dm.add_argument('--ppa', action='append',
                    help='PPA to add. Can be repeated.')
    dm.add_argument('--package-dir', action='append',
                    help='''Directory containing debs to use as an on device
                    repository. Can be repeated.''')
    dm.add_argument('-p', '--package', action='append',
                    help='Package to install. Can be repeated.')
    dm.add_argument('-r', '--remotepassword', action='append',
                    help='Password to use for sudo on the device.')
    dm.set_defaults(func=_handle_writable_image)

    ei = sub.add_parser('edges-intro', help='Enable/disable the edges intro.')
    ei.set_defaults(func=_handle_edges_intro)
    ei = ei.add_mutually_exclusive_group(required=True)
    ei.add_argument('--enable', action='store_true', required=False,
                    help='enable edges intro')
    ei.add_argument('--disable', action='store_false', required=False,
                    help='disable edges intro')

    ei = sub.add_parser('welcome-wizard',
                        help='Enable/disable the welcome wizard.')
    ei.set_defaults(func=_handle_welcome_wizard)
    ei = ei.add_mutually_exclusive_group(required=True)
    ei.add_argument('--enable', action='store_true', required=False,
                    help='enable welcome wizard')
    ei.add_argument('--disable', action='store_false', required=False,
                    help='disable welcome wizard')

    nw = sub.add_parser('network',
                        help='Read/write current wifi connection.')
    nw.set_defaults(func=_handle_network)
    nw_grp = nw.add_mutually_exclusive_group(required=True)
    nw_grp.add_argument('-r', '--read', action='store_true',
                        help='Read current wifi ESSID and key in a format '
                        'suitable for --write')
    nw_grp.add_argument('-w', '--write', metavar='WIFI_SPEC',
                        help='Connect to given wifi ESSID and key (format '
                        'from --read)')
    return parser.parse_args()


def main():
    args = parse_arguments()
    adb = UbuntuDevice(args.serial)
    adb.start()
    args.func(adb, args)

if __name__ == '__main__':
    main()