crdev

#!/usr/bin/env python
# Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This script provides a suite of utility functions for the chroot.
"""

import optparse
import os
import shutil
import subprocess
import sys

_PROG = "crdev"
_SSH_KEY_FILEPATH = os.path.expanduser("~/.ssh/id_rsa")
_PUB_KEY_FILENAME = "~/.ssh/id_rsa.pub"
_PUB_KEY_FILEPATH = os.path.expanduser(_PUB_KEY_FILENAME)
_AUTH_KEYS_FILEPATH = os.path.expanduser("~/.ssh/authorized_keys")
_CHROME_MOUNT_PATH = "/tmp/chrome"
_DEFAULT_CHROME_DIR = "~/chromeos/chromium"

RAN_SUDO = False

###############################################################################
# Utility commands

def _Confirm(message):
  """Request confirmation and return True/False."""
  reply = raw_input(message + " [y/N] ")
  if reply == "y" or reply == "Y":
    return True
  return False


def _CheckOverwriteFile(filepath):
  """Check for a file and request confirmation if it exists."""
  if os.path.isfile(filepath):
    if not _Confirm(filepath + " already exists. Overwrite?"):
      return False
  return True


def _WriteFile(filepath, contents):
  """Write string |contents| to |filepath|."""
  try:
    print "Writing to: ", filepath
    with open(filepath,"w") as f_1:
      f_1.write(contents)
  except EnvironmentError:
    print >> sys.stderr, "Failed to write to file: ", filepath
    return False

  return True


def _Sudo():
  """Request sudo access with message."""
  global RAN_SUDO
  if not RAN_SUDO:
    print "Executing: sudo -v"
  try:
    subprocess.call(["sudo", "-v"])
  except EnvironmentError:
    print >> sys.stderr, "Failed to run sudo"
    return False
  RAN_SUDO = True
  return True


def _RunCommand(args):
  """Pass |args| to subprocess.call() and check the result."""
  cmd = ''.join([x + " " for x in args])
  try:
    if args[0] == "sudo":
      if _Sudo() == False:
        return False
    print "Executing: ", cmd
    retcode = subprocess.call(args)
  except EnvironmentError:
    print >> sys.stderr, "Failed to run command: ", cmd
    return False
  else:
    if retcode != 0:
      print >> sys.stderr, "Error: ", retcode, ", running command: ", cmd
      return False
    return True


def _GetCommandOutput(args):
  """Pass |args| to subprocess.Popen and return the output."""
  cmd = ''.join([x + " " for x in args])
  try:
    proc = subprocess.Popen(args, stdout=subprocess.PIPE)
    result = proc.communicate()[0]
  except EnvironmentError:
    print >> sys.stderr, "Failed to run command: ", cmd
    return None
  result = result.rstrip('\n')
  return result


def _MakeDirectory(dirpath):
  """Create directory |dirpath| if it does not exist."""
  if not os.path.isdir(dirpath):
    try:
      print "Creating directory: ", dirpath
      os.makedirs(dirpath)
    except EnvironmentError:
      print >> sys.stderr, "Failed to create directory: ", dirpath
      return False
  return True


def _RemoveFile(filepath):
  """Remove file |filepath| if it exists."""
  if os.path.isfile(filepath):
    try:
      print "Removing file: ", filepath
      os.remove(filepath)
    except EnvironmentError:
      print >> sys.stderr, "Failed to remove file: ", filepath
      return False
  return True


def _RemoveDirectory(dirpath):
  """Recursively remove directory |dirpath| if it exists."""
  if os.path.isdir(dirpath):
    try:
      print "Removing directory: ", dirpath
      shutil.rmtree(dirpath)
    except EnvironmentError:
      print >> sys.stderr, "Failed to remove directory: ", dirpath
      return False
  return True


def _DevUser():
  """Extract the user name from lsb-release."""
  awk_expr  = """/CHROMEOS_RELEASE_DESCRIPTION/ {"""
  awk_expr += """ sub(/.*Developer Build - /,"");"""
  awk_expr += """ sub(/\).*/,"");"""
  awk_expr += """ print; }"""
  return _GetCommandOutput(["awk", awk_expr, "/etc/lsb-release"])


def _DevHost():
  """Extract the host name from lsb-release."""
  awk_expr  = """/CHROMEOS_DEVSERVER/ {"""
  awk_expr += """ sub(/.*http:\/\//,"");"""
  awk_expr += """ sub(/:8080.*/,"");"""
  awk_expr += """ print; }"""
  return _GetCommandOutput(["awk", awk_expr, "/etc/lsb-release"])

###############################################################################
# Other Commands

def TestCommand(args):
  """Test command."""
  if len(args) == 0:
    args = ["sudo", "ls", "/"]
    return _RunCommand(args)
  return False


def GetUser(unused_args=0):
  """Gets the user name from /etc/lsb-release."""
  print _DevUser()
  return True


def GetHost(unused_args=0):
  """Gets the host name from /etc/lsb-release."""
  print _DevHost()
  return True


def MountWriteable(unused_args=0):
  """Remounts / as rw."""
  return _RunCommand(["sudo", "mount", "-o", "remount,rw", "/"])


def ShowIP(unused_args=0):
  """Shows the IP address of the device."""
  proc1 = subprocess.Popen(["/sbin/ifconfig", "eth0"], stdout=subprocess.PIPE)
  awk_cmd  = """/inet addr/ {"""
  awk_cmd += """ sub(/.*inet addr:/,""); sub(/  Bcast:.*/,"");"""
  awk_cmd += """ print; }"""
  proc2 = subprocess.Popen(
      ["awk", awk_cmd], stdin=proc1.stdout, stdout=subprocess.PIPE)
  proc1.stdout.close()
  result = proc2.communicate()[0].rstrip('\n')
  print "IP: ", result
  return True


def KillChrome(unused_args=0):
  """Kills all chrome processes and prevents restarting of chrome."""
  res = True
  res &=  _RunCommand(["touch", "/tmp/disable_chrome_restart"])
  res &= _RunCommand(["sudo", "pkill", "-9", "chrome"])
  return res


def ShowOobe(unused_args=0):
  """Removes .oobe_completed and Local State directory."""

  res = True
  res &= KillChrome()
  res &= _RemoveFile("/home/chronos/.oobe_completed")
  res &= _RemoveDirectory("/home/chronos/Local State")

  if not res:
    print >> sys.stderr, "Unable to set up OOBE mode."
    return False

  if _Confirm("Reboot is required to enter OOBE mode. Reboot now?"):
    _RunCommand(["sudo", "reboot"])
  return True

###############################################################################
# Setup Commands

def SetBash(unused_args):
  """Sets the default shell to bash."""
  if not _Sudo():
    return False
  if not MountWriteable():
    return False
  res = True
  res &= _RunCommand(["chsh", "-s", "/bin/bash"])
  res &= _RunCommand(["chsh", "-s", "/bin/bash", "chronos"])
  return res


def SetupBashrc(unused_args):
  """Sets up .bashrc."""
  filepath = os.path.expanduser("~/.bashrc")
  if not _CheckOverwriteFile(filepath):
    return False

  print "Writing to: ", filepath
  bashrc  = "#!/bin/bash\n"
  bashrc += "# .bashrc file set by %s\n" % _PROG
  bashrc += "export DISPLAY=:0.0\n"
  bashrc += "export PATH=$PATH:/sbin:/usr/sbin:/usr/local/sbin\n"
  bashrc += "/sbin/ifconfig eth0 | grep 'inet addr'\n"
  if not _WriteFile(filepath, bashrc):
    return False

  filepath = os.path.expanduser("~/.bash_profile")
  print "Writing to: ", filepath
  bashprofile  = "#!/bin/bash\n"
  bashprofile += ". $HOME/.bashrc\n"
  if not _WriteFile(filepath, bashprofile):
    return False
  return True


def SetupDev(unused_args):
  """Developer friendly setup: skip oobe, don't forget usernames,
\t\tdisable suspend, enable ssh and remote debugging,
\t\trun commands from /tmp,/home."""
  if not _Sudo():
    return False
  if not MountWriteable():
    return False
  res = True
  res &= _WriteFile("1", "/home/chronos/.oobe_completed")
  res &= _RemoveFile("/root/.forget_usernames")
  res &= _MakeDirectory("/usr/share/power_manager")
  res &= _WriteFile("1", "/usr/share/power_manager/disable_idle_suspend")
  # Enable iptables and system-services for remote debugging
  for filename in ["iptables", "saft"]:
    res &= _RunCommand(["sudo", "sed", "-i", "-e",
                       "s/#for_test //", "/etc/init/%s.conf" % filename])
  # Allow commands to be run from /home and /tmp
  res &= _RunCommand(["sudo", "sed", "-i", "-e",
                     "s/#mod_for_test#//g", "/sbin/chromeos_startup"])
  return res


def SetupSsh(unused_args):
  """Sets up ssh configuration so that the dev host can ssh to the device."""
  if not MountWriteable():
    return False

  user = _DevUser()
  host = _DevHost()

  if not _CheckOverwriteFile(_SSH_KEY_FILEPATH):
    return False

  res = True
  res &= _RemoveFile(_SSH_KEY_FILEPATH)
  # Enable ssh to device
  res &= _RunCommand(
      ["sudo", "sed", "-i", "s/#for_test //", "/etc/init/openssh-server.conf"])
  # Generate an ssh key
  if _RunCommand(["ssh-keygen", "-f", _SSH_KEY_FILEPATH, "-N", "", "-q"]):
    host_source_path = "%s@%s:%s" % (user, host, _PUB_KEY_FILENAME)
    # Copy the ssh key to the host
    res &= _RunCommand(["scp", host_source_path, _AUTH_KEYS_FILEPATH])

  return res


def AuthorizeSsh(unused_args):
  """Authorizes this netbook to connect to your dev host.
\t\t*Only use this on a device in a secure location!*"""

  user = _DevUser()
  host = _DevHost()

  if not _Confirm("This will copy %s to authorized_keys on %s. Are you sure?" %
                  (_PUB_KEY_FILENAME, host)):
    return False

  proc1 = subprocess.Popen(["cat", _PUB_KEY_FILEPATH], stdout=subprocess.PIPE)
  try:
    ssh_args = ["ssh", user+"@"+host, "cat >> ~/.ssh/authorized_keys"]
    proc2 = subprocess.Popen(
        ssh_args, stdin=proc1.stdout, stdout=subprocess.PIPE)
  except EnvironmentError:
    print >> sys.stderr, "Error executing: ", ' '.join(ssh_args)
    return False
  try:
    proc1.stdout.close()
    result, err = proc2.communicate()
  except EnvironmentError:
    print >> sys.stderr, \
        "Error completing ssh command: ", result, "Error: ", err
    return False
  return True


def SetupSshfsForChrome(unused_args):
  """<chrome-drir> Sets up sshfs mount to chrome directory on host."""

  user = _DevUser()
  host = _DevHost()

  prompt = "Chrome directory on %s [ %s ]: " % (host, _DEFAULT_CHROME_DIR)
  chromedir = raw_input(prompt).rstrip()
  if not chromedir:
    chromedir = _DEFAULT_CHROME_DIR

  target = user+"@"+host+":"+chromedir
  print "Setting up sshfs mount to: ", target

  res = True
  res &= _RunCommand(["sudo", "modprobe", "fuse"])
  if os.path.isdir(_CHROME_MOUNT_PATH):
    res &= _RunCommand(["fusemount", "-q", "-u", _CHROME_MOUNT_PATH])
    res &= _RemoveDirectory(_CHROME_MOUNT_PATH)
  res &= _MakeDirectory(_CHROME_MOUNT_PATH)
  res &= _RunCommand(["sshfs", target, _CHROME_MOUNT_PATH])
  res &= _RunCommand(["sudo", "/sbin/iptables",
                     "-A", "INPUT", "-p", "tcp", "--dport", "1234",
                     "-j", "ACCEPT"])
  return res

###############################################################################
# Multi-commands (convenience functions)

def Setup(args):
  """Performs default developer setup (bash,bashrc,dev,ssh)."""
  if not SetBash(args):
    if not _Confirm("Bash setup failed. Continue?"):
      return False
  if not SetupBashrc(args):
    if not _Confirm(".bashrc setup failed. Continue?"):
      return False
  if not SetupDev(args):
    if not _Confirm("Dev setup failed. Continue?"):
      return False
  if not SetupSsh(args):
    return False
  return True

###############################################################################

_SETUP_COMMANDS = {
    'setup': Setup,
    'dev': SetupDev,
    'bash': SetBash,
    'bashrc': SetupBashrc,
    'ssh': SetupSsh,
    'sshauthorize': AuthorizeSsh,
    'sshfs': SetupSshfsForChrome,
}


_OTHER_COMMANDS = {
    'mountrw': MountWriteable,
    'ip': ShowIP,
    'test': TestCommand,
    'user': GetUser,
    'host': GetHost,
}


def GetUsage(commands):
  """Get the docstring for each command."""
  usage = ""
  for cmd in commands.items():
    usage += " "
    usage += cmd[0]
    usage += ":\t"
    if len(cmd[0]) < 6:
      usage += "\t"
    doc = cmd[1].__doc__
    if doc:
      usage += doc
    usage += "\n"
  return usage


###############################################################################

def main():
  """Main crdev function"""
  usage = """usage: crdev command [options]

Note: Beta! Send feature requests / changes to stevenjb@chromium.org (for now)

"""

  usage += "Setup Commands:\n"
  usage += GetUsage(_SETUP_COMMANDS)
  usage += "Other Commands:\n"
  usage += GetUsage(_OTHER_COMMANDS)

  parser = optparse.OptionParser(usage)
  args = parser.parse_args()[1]

  if not args:
    print usage
    return

  if os.getenv('USER') != 'chronos':
    print >> sys.stderr, "%s must be run as chronos." % _PROG
    return

  cmd = args[0]
  if cmd in _SETUP_COMMANDS.keys():
    res = _SETUP_COMMANDS[cmd](args[1:])
  elif cmd in _OTHER_COMMANDS.keys():
    res = _OTHER_COMMANDS[cmd](args[1:])
  else:
    parser.error("Unknown command: " + cmd)
    return

  if not res:
    print >> sys.stderr, "Error running command: " + ' '.join(args)
  else:
    print "Success!"


if __name__ == '__main__':
  main()