#!/usr/bin/env python
#
# Copyright 2015 The Chromium 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 is a script developers can use to set-up their workstation to let
# Telemetry read the CPU's Model Specific Registers in order to get power
# measurements. It can check if reading from MSRs is possible as any user, but
# must run as root to make changes. Not all changes are sticky, so one has to
# re-run this script after each reboot.
#
# This script is currently Debian/Ubuntu specific.

from __future__ import print_function
from __future__ import absolute_import
import os
import subprocess
import sys

MSR_DEV_FILE_PATH = '/dev/cpu/0/msr'
RDMSR_PATH = '/usr/sbin/rdmsr'

def _Usage(prog_name):
  """Print a help message."""
  print('Run "%s" as a regular user to check if reading from the MSR ' \
      'is possible.' % prog_name)
  print('Run "%s enable" as root to automatically set up reading from ' \
      'the MSR.' % prog_name)


def _CheckMsrKernelModule():
  """Return whether the 'msr' kernel module is loaded."""
  proc = subprocess.Popen('/sbin/lsmod', stdout=subprocess.PIPE)
  stdout = proc.communicate()[0]
  ret = proc.wait()
  if ret != 0:
    raise OSError('lsmod failed')

  if not any(line.startswith('msr ') for line in stdout.splitlines()):
    print('Error: MSR module not loaded.')
    return False

  return True


def _CheckMsrDevNodes():
  """Check whether the MSR /dev files have the right permissions."""
  if not os.path.exists(MSR_DEV_FILE_PATH):
    print('Error: %s does not exist.' % MSR_DEV_FILE_PATH)
    return False

  if not os.access(MSR_DEV_FILE_PATH, os.R_OK):
    print('Error: Cannot read from %s' % MSR_DEV_FILE_PATH)
    return False

  return True


def _CheckRdmsr():
  """Check and make sure /usr/sbin/rdmsr is set up correctly."""
  if not os.access(RDMSR_PATH, os.X_OK):
    print('Error: %s missing or not executable.' % RDMSR_PATH)
    return False

  proc = subprocess.Popen(['/sbin/getcap', RDMSR_PATH], stdout=subprocess.PIPE)
  stdout = proc.communicate()[0]
  ret = proc.wait()
  if ret != 0:
    raise OSError('getcap failed')

  if 'cap_sys_rawio+ep' not in stdout:
    print('Error: /usr/sbin/rdmsr needs RAWIO capability.')
    return False

  return True


def _RunAllChecks():
  """Check to make sure it is possible to read from the MSRs."""
  if os.geteuid() == 0:
    print('WARNING: Running as root, msr permission check likely inaccurate.')

  has_dev_node = _CheckMsrDevNodes() if _CheckMsrKernelModule() else False
  has_rdmsr = _CheckRdmsr()
  return has_dev_node and has_rdmsr


def _EnableMsr(prog_name):
  """Do all the setup needed to pass _RunAllChecks().

  Needs to run as root."""
  if os.geteuid() != 0:
    print('Error: Must run "%s enable" as root.' % prog_name)
    return False

  print('Loading msr kernel module.')
  ret = subprocess.call(['/sbin/modprobe', 'msr'])
  if ret != 0:
    print('Error: Cannot load msr module.')
    return False

  print('Running chmod on %s.' % MSR_DEV_FILE_PATH)
  ret = subprocess.call(['/bin/chmod', 'a+r', MSR_DEV_FILE_PATH])
  if ret != 0:
    print('Error: Cannot chmod %s.' % MSR_DEV_FILE_PATH)
    return False

  if not os.access(RDMSR_PATH, os.F_OK):
    print('Need to install the msr-tools package.')
    ret = subprocess.call(['/usr/bin/apt-get', 'install', '-y', 'msr-tools'])
    if ret != 0:
      print('Error: Did not successfully install msr-tools.')
      return False

  print('Running setcap on %s.' % RDMSR_PATH)
  ret = subprocess.call(['/sbin/setcap', 'cap_sys_rawio+ep', RDMSR_PATH])
  if ret != 0:
    print('Error: Cannot give /usr/sbin/rdmsr RAWIO capability.')
    return False

  return True


def main(prog_name, argv):
  if len(argv) == 0:
    if _RunAllChecks():
      print('Check succeeded')
      return 0

    print('Check failed, try running "%s enable" as root to fix.' % prog_name)
    return 1

  if len(argv) == 1:
    if argv[0] == 'enable':
      return 0 if _EnableMsr(prog_name) else 1

    print('Error: Unknown sub-command %s' % argv[0])
    _Usage(prog_name)
    return 1

  print('Error: Bad number of arguments')
  _Usage(prog_name)
  return 1


if __name__ == '__main__':
  sys.exit(main(os.path.basename(sys.argv[0]), sys.argv[1:]))