go/src/infra/tools/kitchen/testdata/recipe_repo/recipes

#!/usr/bin/env python
# This file mocks typical recipes.py that normally runs a recipe.

import argparse
import json
import os
import shutil
import sys
import urllib2


def main():
  parser = argparse.ArgumentParser()

  subparsers = parser.add_subparsers()
  fetch_cmd = subparsers.add_parser('fetch')
  fetch_cmd.set_defaults(command='fetch')

  run_cmd = subparsers.add_parser('run')
  run_cmd.add_argument('--output-result-json')
  run_cmd.add_argument('--properties-file')
  run_cmd.set_defaults(command='run')

  args, _ = parser.parse_known_args()

  # See cook_test.go
  assert os.environ.get('SWARMING_TASK_ID') == 'task'

  if args.command == 'fetch':
    # Fetch happens under the system account. See localauth.Server config in
    # cook_test.go.
    assert get_current_account() == 'system_acc', get_current_account()
    assert get_git_email() == 'system@example.com', get_git_email()
    return 0

  assert args.command == 'run'
  assert args.output_result_json
  assert args.properties_file

  # Actual recipe execution happens under the recipe account. See
  # localauth.Server config in cook_test.go.
  assert get_current_account() == 'recipe_acc', get_current_account()
  assert get_git_email() == 'recipe@example.com', get_git_email()
  assert get_metadata_token() == 'fake_token_0', get_metadata_token()

  with open(args.properties_file) as f:
    properties = json.load(f)
  cfg = properties.pop('recipe_mock_cfg')

  with open(cfg['input_path'], 'w') as f:
    json.dump({
      'args': sys.argv,
      'properties': properties,
    }, f)

  mocked_result_path = cfg.get('mocked_result_path')
  if mocked_result_path:
    shutil.copyfile(mocked_result_path, args.output_result_json)
  return cfg['exitCode']


def get_current_account():
  with open(os.environ['LUCI_CONTEXT'], 'rt') as f:
    lc = json.load(f)
  return lc["local_auth"]["default_account_id"]


def get_git_email():
  home = os.environ['INFRA_GIT_WRAPPER_HOME']
  with open(os.path.join(home, '.gitconfig'), 'rt') as f:
    cfg = f.read()
  for line in cfg.splitlines():
    line = line.strip()
    if line.startswith('email = '):
      return line[len('email = '):]
  return None


def get_metadata_token():
  metasrv = os.environ.get('GCE_METADATA_HOST')
  assert metasrv and metasrv.startswith('127.0.0.1:'), metasrv
  req = urllib2.Request(
      'http://' + metasrv + '/computeMetadata/v1/instance/'
          'service-accounts/default/token',
      headers={'Metadata-Flavor': 'Google'})
  resp = json.loads(urllib2.urlopen(req).read())
  return resp['access_token']


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