<reponame>xswz8015/infra
# Copyright 2021 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.

import subprocess
import argparse
import sys
import os
import re
import json
import codecs
import platform

# Used to run commands through powershell on windows and collect the responses
# and log files that may be generated by the command run. This is meant for use
# in situations where the commands either don't return the results to STDOUT,
# instead writing to a log file, or if the commands do return results but do not
# write the logs to STDOUT. This can also be used to run regular windows
# executables as well, powershell will happily execute them for you.

# ALLOWLIST filters the files that can be read to STDOUT
ALLOWLIST = [
    re.compile(y) for y in [
        '^.*\.log$',  # Log files
    ]
]

# TODO(anushruth): Update list with all other possibilities.
codec_map = {
    codecs.BOM_UTF16_LE: 'utf-16-le',
    codecs.BOM_UTF32: 'utf-32',
    codecs.BOM_UTF8: 'utf-8-sig'
}


def main(argv):
  """Runs the given powershell command and writes all the logs to stdout """
  # Ensure we are running on windows. Until MS releases powershell for other
  # platforms
  if platform.system() != 'Windows':
    print(json.dumps(gen_result('Not run on Windows')))
    return
  parser = argparse.ArgumentParser(
      description="""Runs a powershell command,
                     waits for its completion and returns all
                     the output logs to stdout as a json""",
      epilog="""Meant to be used by the powershell
                recipe module to run a command that
                generates multiple logs and stream the logs
                back to the recipe""")

  # command is the actual powershell command required to be run
  parser.add_argument(
      '--command', required=True, help='powershell command to execute')
  parser.add_argument(
      '--logs',
      required=False,
      nargs='*',
      help='log file or dir to watch and return stuff from')
  parser.add_argument(
      '--ret_codes',
      required=False,
      type=int,
      default=[0],
      nargs='*',
      help='return codes to treat as success')
  parser.add_argument(
      'args', nargs='*', help='optionals args to the powershell command')

  iput = parser.parse_args(argv[1:])
  logs = exec_ps(
      iput.command, iput.logs, args=iput.args, ret_codes=iput.ret_codes)
  print(json.dumps(logs))


def ensure_logs(logs):
  """ Checks if any log dir doesn't exist and creates it"""
  if not logs:
    return
  for log in logs:
    if not os.path.exists(log) and not is_allowed(log):
      # If the path doesn't exist and is not a file. Create the dir
      os.makedirs(log)


def exec_ps(command, logs, args, ret_codes):
  """ Executes a power shell command and waits for it to complete.
      Returns all the logs on completion as a json to stdout.

      command: path to script/executable/batchfile, powershell command
      logs: list of log files and directories.
      args: optional args to the command
      ret_codes: optional success return codes

      Returns dict containing keys 'results' and every '<log-file>' in logs."""

  ensure_logs(logs)
  # powershell command
  psc = ['powershell', '-Command', command] + args
  # l contains all the logs + return values
  l = {}
  # Attempt to execute the command
  try:
    output = subprocess.check_output(psc, stderr=subprocess.STDOUT)

    try:
      # Check if the output is a json file
      jout = json.loads(output)
    except Exception as j:
      # It's not known if the script completed successfully
      l = gen_result(
          'No json object returned. Check stdout_stderr for ' +
          'results. {}'.format(j), True)
      # not a json return script
      l['stdout_stderr'] = output
    else:
      # It's a json file
      l['results'] = jout

  except subprocess.CalledProcessError as e:
    # Non zero return by the script/cmd run
    l = gen_result(e.output, e.returncode in ret_codes)
    l['results']['Command'] = e.cmd
    l['results']['ReturnCode'] = e.returncode
    l['stdout_stderr'] = e.output

  except Exception as e:
    # Failed to run the command for some reason
    l = gen_result(str(e))
    l['results']['Command'] = ' '.join(psc)

  finally:
    # Read all the logs to stdout
    if logs:
      for log in logs:
        if os.path.isdir(log):
          for k, v in read_logs([os.path.join(log, x) for x in os.listdir(log)
                                ]).items():
            l[k] = v
        else:
          op = read_logs(log, l)
          l[log] = op[log]

  return l


def read_logs(logs):
  """ Reads all the given files to RAM and returns a dict of contents.

      logs: list of log files and directories.

      Returns dict containing keys for each log file and its contents
      as value. """

  l = {}
  for log in logs:
    if not os.path.isdir(log) and is_allowed(log):
      f = open(log, 'r')
      contents = f.read()
      l[log] = contents
      # Some logs may be encoded differently. Check if they have the unicode
      # BOM at the start and decode them. ADKSetup is known to generate logs
      # with different encodings during the same run.
      for k, v in codec_map.items():
        if len(contents) >= len(k) and k == contents[:len(k)]:
          # See codec_map for codec to decode string mapping
          l[log] = contents.decode(v)
          break
      f.close()
  return l


def is_allowed(l):
  """ Implements ALLOWLIST

      l: log file to check

      Returns True if l matches anything in ALLOWLIST, false otherwise. """

  for d in ALLOWLIST:
    if d.match(l) != None:
      return True
  return False


def gen_result(err, success=False):
  """ gen_result returns the result dict with given params"""
  return {'results': {'Success': success, 'ErrorInfo': {'Message': err,}}}


if __name__ == '__main__':
  main(sys.argv)

# Copyright 2021 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.

import subprocess
import argparse
import sys
import os
import re
import json
import codecs
import platform

# Used to run commands through powershell on windows and collect the responses
# and log files that may be generated by the command run. This is meant for use
# in situations where the commands either don't return the results to STDOUT,
# instead writing to a log file, or if the commands do return results but do not
# write the logs to STDOUT. This can also be used to run regular windows
# executables as well, powershell will happily execute them for you.

# ALLOWLIST filters the files that can be read to STDOUT
ALLOWLIST = [
    re.compile(y) for y in [
        '^.*\.log$',  # Log files
    ]
]

# TODO(anushruth): Update list with all other possibilities.
codec_map = {
    codecs.BOM_UTF16_LE: 'utf-16-le',
    codecs.BOM_UTF32: 'utf-32',
    codecs.BOM_UTF8: 'utf-8-sig'
}


def main(argv):
  """Runs the given powershell command and writes all the logs to stdout """
  # Ensure we are running on windows. Until MS releases powershell for other
  # platforms
  if platform.system() != 'Windows':
    print(json.dumps(gen_result('Not run on Windows')))
    return
  parser = argparse.ArgumentParser(
      description="""Runs a powershell command,
                     waits for its completion and returns all
                     the output logs to stdout as a json""",
      epilog="""Meant to be used by the powershell
                recipe module to run a command that
                generates multiple logs and stream the logs
                back to the recipe""")

  # command is the actual powershell command required to be run
  parser.add_argument(
      '--command', required=True, help='powershell command to execute')
  parser.add_argument(
      '--logs',
      required=False,
      nargs='*',
      help='log file or dir to watch and return stuff from')
  parser.add_argument(
      '--ret_codes',
      required=False,
      type=int,
      default=[0],
      nargs='*',
      help='return codes to treat as success')
  parser.add_argument(
      'args', nargs='*', help='optionals args to the powershell command')

  iput = parser.parse_args(argv[1:])
  logs = exec_ps(
      iput.command, iput.logs, args=iput.args, ret_codes=iput.ret_codes)
  print(json.dumps(logs))


def ensure_logs(logs):
  """ Checks if any log dir doesn't exist and creates it"""
  if not logs:
    return
  for log in logs:
    if not os.path.exists(log) and not is_allowed(log):
      # If the path doesn't exist and is not a file. Create the dir
      os.makedirs(log)


def exec_ps(command, logs, args, ret_codes):
  """ Executes a power shell command and waits for it to complete.
      Returns all the logs on completion as a json to stdout.

      command: path to script/executable/batchfile, powershell command
      logs: list of log files and directories.
      args: optional args to the command
      ret_codes: optional success return codes

      Returns dict containing keys 'results' and every '<log-file>' in logs."""

  ensure_logs(logs)
  # powershell command
  psc = ['powershell', '-Command', command] + args
  # l contains all the logs + return values
  l = {}
  # Attempt to execute the command
  try:
    output = subprocess.check_output(psc, stderr=subprocess.STDOUT)

    try:
      # Check if the output is a json file
      jout = json.loads(output)
    except Exception as j:
      # It's not known if the script completed successfully
      l = gen_result(
          'No json object returned. Check stdout_stderr for ' +
          'results. {}'.format(j), True)
      # not a json return script
      l['stdout_stderr'] = output
    else:
      # It's a json file
      l['results'] = jout

  except subprocess.CalledProcessError as e:
    # Non zero return by the script/cmd run
    l = gen_result(e.output, e.returncode in ret_codes)
    l['results']['Command'] = e.cmd
    l['results']['ReturnCode'] = e.returncode
    l['stdout_stderr'] = e.output

  except Exception as e:
    # Failed to run the command for some reason
    l = gen_result(str(e))
    l['results']['Command'] = ' '.join(psc)

  finally:
    # Read all the logs to stdout
    if logs:
      for log in logs:
        if os.path.isdir(log):
          for k, v in read_logs([os.path.join(log, x) for x in os.listdir(log)
                                ]).items():
            l[k] = v
        else:
          op = read_logs(log, l)
          l[log] = op[log]

  return l


def read_logs(logs):
  """ Reads all the given files to RAM and returns a dict of contents.

      logs: list of log files and directories.

      Returns dict containing keys for each log file and its contents
      as value. """

  l = {}
  for log in logs:
    if not os.path.isdir(log) and is_allowed(log):
      f = open(log, 'r')
      contents = f.read()
      l[log] = contents
      # Some logs may be encoded differently. Check if they have the unicode
      # BOM at the start and decode them. ADKSetup is known to generate logs
      # with different encodings during the same run.
      for k, v in codec_map.items():
        if len(contents) >= len(k) and k == contents[:len(k)]:
          # See codec_map for codec to decode string mapping
          l[log] = contents.decode(v)
          break
      f.close()
  return l


def is_allowed(l):
  """ Implements ALLOWLIST

      l: log file to check

      Returns True if l matches anything in ALLOWLIST, false otherwise. """

  for d in ALLOWLIST:
    if d.match(l) != None:
      return True
  return False


def gen_result(err, success=False):
  """ gen_result returns the result dict with given params"""
  return {'results': {'Success': success, 'ErrorInfo': {'Message': err,}}}


if __name__ == '__main__':
  main(sys.argv)
# Copyright 2021 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. # Used to run commands through powershell on windows and collect the responses # and log files that may be generated by the command run. This is meant for use # in situations where the commands either don't return the results to STDOUT, # instead writing to a log file, or if the commands do return results but do not # write the logs to STDOUT. This can also be used to run regular windows # executables as well, powershell will happily execute them for you. # ALLOWLIST filters the files that can be read to STDOUT # Log files # TODO(anushruth): Update list with all other possibilities. Runs the given powershell command and writes all the logs to stdout # Ensure we are running on windows. Until MS releases powershell for other # platforms Runs a powershell command,                      waits for its completion and returns all                      the output logs to stdout as a json Meant to be used by the powershell                 recipe module to run a command that                 generates multiple logs and stream the logs                 back to the recipe # command is the actual powershell command required to be run Checks if any log dir doesn't exist and creates it # If the path doesn't exist and is not a file. Create the dir Executes a power shell command and waits for it to complete.       Returns all the logs on completion as a json to stdout.        command: path to script/executable/batchfile, powershell command       logs: list of log files and directories.       args: optional args to the command       ret_codes: optional success return codes        Returns dict containing keys 'results' and every '<log-file>' in logs. # powershell command # l contains all the logs + return values # Attempt to execute the command # Check if the output is a json file # It's not known if the script completed successfully # not a json return script # It's a json file # Non zero return by the script/cmd run # Failed to run the command for some reason # Read all the logs to stdout Reads all the given files to RAM and returns a dict of contents.        logs: list of log files and directories.        Returns dict containing keys for each log file and its contents       as value. # Some logs may be encoded differently. Check if they have the unicode # BOM at the start and decode them. ADKSetup is known to generate logs # with different encodings during the same run. # See codec_map for codec to decode string mapping Implements ALLOWLIST        l: log file to check        Returns True if l matches anything in ALLOWLIST, false otherwise. gen_result returns the result dict with given params