1*8975f5c5SAndroid Build Coastguard Worker#!/usr/bin/env python3 2*8975f5c5SAndroid Build Coastguard Worker# Copyright 2019 The Chromium Authors 3*8975f5c5SAndroid Build Coastguard Worker# Use of this source code is governed by a BSD-style license that can be 4*8975f5c5SAndroid Build Coastguard Worker# found in the LICENSE file. 5*8975f5c5SAndroid Build Coastguard Worker 6*8975f5c5SAndroid Build Coastguard Worker"""Wraps an executable and any provided arguments into an executable script.""" 7*8975f5c5SAndroid Build Coastguard Worker 8*8975f5c5SAndroid Build Coastguard Workerimport argparse 9*8975f5c5SAndroid Build Coastguard Workerimport os 10*8975f5c5SAndroid Build Coastguard Workerimport sys 11*8975f5c5SAndroid Build Coastguard Workerimport textwrap 12*8975f5c5SAndroid Build Coastguard Worker 13*8975f5c5SAndroid Build Coastguard Worker 14*8975f5c5SAndroid Build Coastguard Worker# The bash template passes the python script into vpython via stdin. 15*8975f5c5SAndroid Build Coastguard Worker# The interpreter doesn't know about the script, so we have bash 16*8975f5c5SAndroid Build Coastguard Worker# inject the script location. 17*8975f5c5SAndroid Build Coastguard WorkerBASH_TEMPLATE = textwrap.dedent("""\ 18*8975f5c5SAndroid Build Coastguard Worker #!/usr/bin/env vpython3 19*8975f5c5SAndroid Build Coastguard Worker _SCRIPT_LOCATION = __file__ 20*8975f5c5SAndroid Build Coastguard Worker {script} 21*8975f5c5SAndroid Build Coastguard Worker """) 22*8975f5c5SAndroid Build Coastguard Worker 23*8975f5c5SAndroid Build Coastguard Worker 24*8975f5c5SAndroid Build Coastguard Worker# The batch template reruns the batch script with vpython, with the -x 25*8975f5c5SAndroid Build Coastguard Worker# flag instructing the interpreter to ignore the first line. The interpreter 26*8975f5c5SAndroid Build Coastguard Worker# knows about the (batch) script in this case, so it can get the file location 27*8975f5c5SAndroid Build Coastguard Worker# directly. 28*8975f5c5SAndroid Build Coastguard WorkerBATCH_TEMPLATE = textwrap.dedent("""\ 29*8975f5c5SAndroid Build Coastguard Worker @SETLOCAL ENABLEDELAYEDEXPANSION \ 30*8975f5c5SAndroid Build Coastguard Worker & vpython3.bat -x "%~f0" %* \ 31*8975f5c5SAndroid Build Coastguard Worker & EXIT /B !ERRORLEVEL! 32*8975f5c5SAndroid Build Coastguard Worker _SCRIPT_LOCATION = __file__ 33*8975f5c5SAndroid Build Coastguard Worker {script} 34*8975f5c5SAndroid Build Coastguard Worker """) 35*8975f5c5SAndroid Build Coastguard Worker 36*8975f5c5SAndroid Build Coastguard Worker 37*8975f5c5SAndroid Build Coastguard WorkerSCRIPT_TEMPLATES = { 38*8975f5c5SAndroid Build Coastguard Worker 'bash': BASH_TEMPLATE, 39*8975f5c5SAndroid Build Coastguard Worker 'batch': BATCH_TEMPLATE, 40*8975f5c5SAndroid Build Coastguard Worker} 41*8975f5c5SAndroid Build Coastguard Worker 42*8975f5c5SAndroid Build Coastguard Worker 43*8975f5c5SAndroid Build Coastguard WorkerPY_TEMPLATE = textwrap.dedent(r""" 44*8975f5c5SAndroid Build Coastguard Worker import os 45*8975f5c5SAndroid Build Coastguard Worker import re 46*8975f5c5SAndroid Build Coastguard Worker import shlex 47*8975f5c5SAndroid Build Coastguard Worker import signal 48*8975f5c5SAndroid Build Coastguard Worker import subprocess 49*8975f5c5SAndroid Build Coastguard Worker import sys 50*8975f5c5SAndroid Build Coastguard Worker import time 51*8975f5c5SAndroid Build Coastguard Worker 52*8975f5c5SAndroid Build Coastguard Worker _WRAPPED_PATH_RE = re.compile(r'@WrappedPath\(([^)]+)\)') 53*8975f5c5SAndroid Build Coastguard Worker _PATH_TO_OUTPUT_DIR = '{path_to_output_dir}' 54*8975f5c5SAndroid Build Coastguard Worker _SCRIPT_DIR = os.path.dirname(os.path.realpath(_SCRIPT_LOCATION)) 55*8975f5c5SAndroid Build Coastguard Worker 56*8975f5c5SAndroid Build Coastguard Worker 57*8975f5c5SAndroid Build Coastguard Worker def ExpandWrappedPath(arg): 58*8975f5c5SAndroid Build Coastguard Worker m = _WRAPPED_PATH_RE.search(arg) 59*8975f5c5SAndroid Build Coastguard Worker if m: 60*8975f5c5SAndroid Build Coastguard Worker head = arg[:m.start()] 61*8975f5c5SAndroid Build Coastguard Worker tail = arg[m.end():] 62*8975f5c5SAndroid Build Coastguard Worker relpath = os.path.join( 63*8975f5c5SAndroid Build Coastguard Worker os.path.relpath(_SCRIPT_DIR), _PATH_TO_OUTPUT_DIR, m.group(1)) 64*8975f5c5SAndroid Build Coastguard Worker npath = os.path.normpath(relpath) 65*8975f5c5SAndroid Build Coastguard Worker if os.path.sep not in npath: 66*8975f5c5SAndroid Build Coastguard Worker # If the original path points to something in the current directory, 67*8975f5c5SAndroid Build Coastguard Worker # returning the normalized version of it can be a problem. 68*8975f5c5SAndroid Build Coastguard Worker # normpath() strips off the './' part of the path 69*8975f5c5SAndroid Build Coastguard Worker # ('./foo' becomes 'foo'), which can be a problem if the result 70*8975f5c5SAndroid Build Coastguard Worker # is passed to something like os.execvp(); in that case 71*8975f5c5SAndroid Build Coastguard Worker # osexecvp() will search $PATH for the executable, rather than 72*8975f5c5SAndroid Build Coastguard Worker # just execing the arg directly, and if '.' isn't in $PATH, this 73*8975f5c5SAndroid Build Coastguard Worker # results in an error. 74*8975f5c5SAndroid Build Coastguard Worker # 75*8975f5c5SAndroid Build Coastguard Worker # So, we need to explicitly return './foo' (or '.\\foo' on windows) 76*8975f5c5SAndroid Build Coastguard Worker # instead of 'foo'. 77*8975f5c5SAndroid Build Coastguard Worker # 78*8975f5c5SAndroid Build Coastguard Worker # Hopefully there are no cases where this causes a problem; if 79*8975f5c5SAndroid Build Coastguard Worker # there are, we will either need to change the interface to 80*8975f5c5SAndroid Build Coastguard Worker # WrappedPath() somehow to distinguish between the two, or 81*8975f5c5SAndroid Build Coastguard Worker # somehow ensure that the wrapped executable doesn't hit cases 82*8975f5c5SAndroid Build Coastguard Worker # like this. 83*8975f5c5SAndroid Build Coastguard Worker return head + '.' + os.path.sep + npath + tail 84*8975f5c5SAndroid Build Coastguard Worker return head + npath + tail 85*8975f5c5SAndroid Build Coastguard Worker return arg 86*8975f5c5SAndroid Build Coastguard Worker 87*8975f5c5SAndroid Build Coastguard Worker 88*8975f5c5SAndroid Build Coastguard Worker def ExpandWrappedPaths(args): 89*8975f5c5SAndroid Build Coastguard Worker for i, arg in enumerate(args): 90*8975f5c5SAndroid Build Coastguard Worker args[i] = ExpandWrappedPath(arg) 91*8975f5c5SAndroid Build Coastguard Worker return args 92*8975f5c5SAndroid Build Coastguard Worker 93*8975f5c5SAndroid Build Coastguard Worker 94*8975f5c5SAndroid Build Coastguard Worker def FindIsolatedOutdir(raw_args): 95*8975f5c5SAndroid Build Coastguard Worker outdir = None 96*8975f5c5SAndroid Build Coastguard Worker i = 0 97*8975f5c5SAndroid Build Coastguard Worker remaining_args = [] 98*8975f5c5SAndroid Build Coastguard Worker while i < len(raw_args): 99*8975f5c5SAndroid Build Coastguard Worker if raw_args[i] == '--isolated-outdir' and i < len(raw_args)-1: 100*8975f5c5SAndroid Build Coastguard Worker outdir = raw_args[i+1] 101*8975f5c5SAndroid Build Coastguard Worker i += 2 102*8975f5c5SAndroid Build Coastguard Worker elif raw_args[i].startswith('--isolated-outdir='): 103*8975f5c5SAndroid Build Coastguard Worker outdir = raw_args[i][len('--isolated-outdir='):] 104*8975f5c5SAndroid Build Coastguard Worker i += 1 105*8975f5c5SAndroid Build Coastguard Worker else: 106*8975f5c5SAndroid Build Coastguard Worker remaining_args.append(raw_args[i]) 107*8975f5c5SAndroid Build Coastguard Worker i += 1 108*8975f5c5SAndroid Build Coastguard Worker if not outdir and 'ISOLATED_OUTDIR' in os.environ: 109*8975f5c5SAndroid Build Coastguard Worker outdir = os.environ['ISOLATED_OUTDIR'] 110*8975f5c5SAndroid Build Coastguard Worker return outdir, remaining_args 111*8975f5c5SAndroid Build Coastguard Worker 112*8975f5c5SAndroid Build Coastguard Worker def InsertWrapperScriptArgs(args): 113*8975f5c5SAndroid Build Coastguard Worker if '--wrapper-script-args' in args: 114*8975f5c5SAndroid Build Coastguard Worker idx = args.index('--wrapper-script-args') 115*8975f5c5SAndroid Build Coastguard Worker args.insert(idx + 1, shlex.join(sys.argv)) 116*8975f5c5SAndroid Build Coastguard Worker 117*8975f5c5SAndroid Build Coastguard Worker def FilterIsolatedOutdirBasedArgs(outdir, args): 118*8975f5c5SAndroid Build Coastguard Worker rargs = [] 119*8975f5c5SAndroid Build Coastguard Worker i = 0 120*8975f5c5SAndroid Build Coastguard Worker while i < len(args): 121*8975f5c5SAndroid Build Coastguard Worker if 'ISOLATED_OUTDIR' in args[i]: 122*8975f5c5SAndroid Build Coastguard Worker if outdir: 123*8975f5c5SAndroid Build Coastguard Worker # Rewrite the arg. 124*8975f5c5SAndroid Build Coastguard Worker rargs.append(args[i].replace('${{ISOLATED_OUTDIR}}', 125*8975f5c5SAndroid Build Coastguard Worker outdir).replace( 126*8975f5c5SAndroid Build Coastguard Worker '$ISOLATED_OUTDIR', outdir)) 127*8975f5c5SAndroid Build Coastguard Worker i += 1 128*8975f5c5SAndroid Build Coastguard Worker else: 129*8975f5c5SAndroid Build Coastguard Worker # Simply drop the arg. 130*8975f5c5SAndroid Build Coastguard Worker i += 1 131*8975f5c5SAndroid Build Coastguard Worker elif (not outdir and 132*8975f5c5SAndroid Build Coastguard Worker args[i].startswith('-') and 133*8975f5c5SAndroid Build Coastguard Worker '=' not in args[i] and 134*8975f5c5SAndroid Build Coastguard Worker i < len(args) - 1 and 135*8975f5c5SAndroid Build Coastguard Worker 'ISOLATED_OUTDIR' in args[i+1]): 136*8975f5c5SAndroid Build Coastguard Worker # Parsing this case is ambiguous; if we're given 137*8975f5c5SAndroid Build Coastguard Worker # `--foo $ISOLATED_OUTDIR` we can't tell if $ISOLATED_OUTDIR 138*8975f5c5SAndroid Build Coastguard Worker # is meant to be the value of foo, or if foo takes no argument 139*8975f5c5SAndroid Build Coastguard Worker # and $ISOLATED_OUTDIR is the first positional arg. 140*8975f5c5SAndroid Build Coastguard Worker # 141*8975f5c5SAndroid Build Coastguard Worker # We assume the former will be much more common, and so we 142*8975f5c5SAndroid Build Coastguard Worker # need to drop --foo and $ISOLATED_OUTDIR. 143*8975f5c5SAndroid Build Coastguard Worker i += 2 144*8975f5c5SAndroid Build Coastguard Worker else: 145*8975f5c5SAndroid Build Coastguard Worker rargs.append(args[i]) 146*8975f5c5SAndroid Build Coastguard Worker i += 1 147*8975f5c5SAndroid Build Coastguard Worker return rargs 148*8975f5c5SAndroid Build Coastguard Worker 149*8975f5c5SAndroid Build Coastguard Worker def ForwardSignals(proc): 150*8975f5c5SAndroid Build Coastguard Worker def _sig_handler(sig, _): 151*8975f5c5SAndroid Build Coastguard Worker if proc.poll() is not None: 152*8975f5c5SAndroid Build Coastguard Worker return 153*8975f5c5SAndroid Build Coastguard Worker # SIGBREAK is defined only for win32. 154*8975f5c5SAndroid Build Coastguard Worker # pylint: disable=no-member 155*8975f5c5SAndroid Build Coastguard Worker if sys.platform == 'win32' and sig == signal.SIGBREAK: 156*8975f5c5SAndroid Build Coastguard Worker print("Received signal(%d), sending CTRL_BREAK_EVENT to process %d" % (sig, proc.pid)) 157*8975f5c5SAndroid Build Coastguard Worker proc.send_signal(signal.CTRL_BREAK_EVENT) 158*8975f5c5SAndroid Build Coastguard Worker else: 159*8975f5c5SAndroid Build Coastguard Worker print("Forwarding signal(%d) to process %d" % (sig, proc.pid)) 160*8975f5c5SAndroid Build Coastguard Worker proc.send_signal(sig) 161*8975f5c5SAndroid Build Coastguard Worker # pylint: enable=no-member 162*8975f5c5SAndroid Build Coastguard Worker if sys.platform == 'win32': 163*8975f5c5SAndroid Build Coastguard Worker signal.signal(signal.SIGBREAK, _sig_handler) # pylint: disable=no-member 164*8975f5c5SAndroid Build Coastguard Worker else: 165*8975f5c5SAndroid Build Coastguard Worker signal.signal(signal.SIGTERM, _sig_handler) 166*8975f5c5SAndroid Build Coastguard Worker signal.signal(signal.SIGINT, _sig_handler) 167*8975f5c5SAndroid Build Coastguard Worker 168*8975f5c5SAndroid Build Coastguard Worker def Popen(*args, **kwargs): 169*8975f5c5SAndroid Build Coastguard Worker assert 'creationflags' not in kwargs 170*8975f5c5SAndroid Build Coastguard Worker if sys.platform == 'win32': 171*8975f5c5SAndroid Build Coastguard Worker # Necessary for signal handling. See crbug.com/733612#c6. 172*8975f5c5SAndroid Build Coastguard Worker kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP 173*8975f5c5SAndroid Build Coastguard Worker return subprocess.Popen(*args, **kwargs) 174*8975f5c5SAndroid Build Coastguard Worker 175*8975f5c5SAndroid Build Coastguard Worker def RunCommand(cmd): 176*8975f5c5SAndroid Build Coastguard Worker process = Popen(cmd) 177*8975f5c5SAndroid Build Coastguard Worker ForwardSignals(process) 178*8975f5c5SAndroid Build Coastguard Worker while process.poll() is None: 179*8975f5c5SAndroid Build Coastguard Worker time.sleep(0.1) 180*8975f5c5SAndroid Build Coastguard Worker return process.returncode 181*8975f5c5SAndroid Build Coastguard Worker 182*8975f5c5SAndroid Build Coastguard Worker 183*8975f5c5SAndroid Build Coastguard Worker def main(raw_args): 184*8975f5c5SAndroid Build Coastguard Worker executable_path = ExpandWrappedPath('{executable_path}') 185*8975f5c5SAndroid Build Coastguard Worker outdir, remaining_args = FindIsolatedOutdir(raw_args) 186*8975f5c5SAndroid Build Coastguard Worker args = {executable_args} 187*8975f5c5SAndroid Build Coastguard Worker InsertWrapperScriptArgs(args) 188*8975f5c5SAndroid Build Coastguard Worker args = FilterIsolatedOutdirBasedArgs(outdir, args) 189*8975f5c5SAndroid Build Coastguard Worker executable_args = ExpandWrappedPaths(args) 190*8975f5c5SAndroid Build Coastguard Worker cmd = [executable_path] + executable_args + remaining_args 191*8975f5c5SAndroid Build Coastguard Worker if executable_path.endswith('.py'): 192*8975f5c5SAndroid Build Coastguard Worker cmd = [sys.executable] + cmd 193*8975f5c5SAndroid Build Coastguard Worker return RunCommand(cmd) 194*8975f5c5SAndroid Build Coastguard Worker 195*8975f5c5SAndroid Build Coastguard Worker 196*8975f5c5SAndroid Build Coastguard Worker if __name__ == '__main__': 197*8975f5c5SAndroid Build Coastguard Worker sys.exit(main(sys.argv[1:])) 198*8975f5c5SAndroid Build Coastguard Worker """) 199*8975f5c5SAndroid Build Coastguard Worker 200*8975f5c5SAndroid Build Coastguard Worker 201*8975f5c5SAndroid Build Coastguard Workerdef Wrap(args): 202*8975f5c5SAndroid Build Coastguard Worker """Writes a wrapped script according to the provided arguments. 203*8975f5c5SAndroid Build Coastguard Worker 204*8975f5c5SAndroid Build Coastguard Worker Arguments: 205*8975f5c5SAndroid Build Coastguard Worker args: an argparse.Namespace object containing command-line arguments 206*8975f5c5SAndroid Build Coastguard Worker as parsed by a parser returned by CreateArgumentParser. 207*8975f5c5SAndroid Build Coastguard Worker """ 208*8975f5c5SAndroid Build Coastguard Worker path_to_output_dir = os.path.relpath( 209*8975f5c5SAndroid Build Coastguard Worker args.output_directory, 210*8975f5c5SAndroid Build Coastguard Worker os.path.dirname(args.wrapper_script)) 211*8975f5c5SAndroid Build Coastguard Worker 212*8975f5c5SAndroid Build Coastguard Worker with open(args.wrapper_script, 'w') as wrapper_script: 213*8975f5c5SAndroid Build Coastguard Worker py_contents = PY_TEMPLATE.format( 214*8975f5c5SAndroid Build Coastguard Worker path_to_output_dir=path_to_output_dir, 215*8975f5c5SAndroid Build Coastguard Worker executable_path=str(args.executable), 216*8975f5c5SAndroid Build Coastguard Worker executable_args=str(args.executable_args)) 217*8975f5c5SAndroid Build Coastguard Worker template = SCRIPT_TEMPLATES[args.script_language] 218*8975f5c5SAndroid Build Coastguard Worker wrapper_script.write(template.format(script=py_contents)) 219*8975f5c5SAndroid Build Coastguard Worker os.chmod(args.wrapper_script, 0o750) 220*8975f5c5SAndroid Build Coastguard Worker 221*8975f5c5SAndroid Build Coastguard Worker return 0 222*8975f5c5SAndroid Build Coastguard Worker 223*8975f5c5SAndroid Build Coastguard Worker 224*8975f5c5SAndroid Build Coastguard Workerdef CreateArgumentParser(): 225*8975f5c5SAndroid Build Coastguard Worker """Creates an argparse.ArgumentParser instance.""" 226*8975f5c5SAndroid Build Coastguard Worker parser = argparse.ArgumentParser() 227*8975f5c5SAndroid Build Coastguard Worker parser.add_argument( 228*8975f5c5SAndroid Build Coastguard Worker '--executable', 229*8975f5c5SAndroid Build Coastguard Worker help='Executable to wrap.') 230*8975f5c5SAndroid Build Coastguard Worker parser.add_argument( 231*8975f5c5SAndroid Build Coastguard Worker '--wrapper-script', 232*8975f5c5SAndroid Build Coastguard Worker help='Path to which the wrapper script will be written.') 233*8975f5c5SAndroid Build Coastguard Worker parser.add_argument( 234*8975f5c5SAndroid Build Coastguard Worker '--output-directory', 235*8975f5c5SAndroid Build Coastguard Worker help='Path to the output directory.') 236*8975f5c5SAndroid Build Coastguard Worker parser.add_argument( 237*8975f5c5SAndroid Build Coastguard Worker '--script-language', 238*8975f5c5SAndroid Build Coastguard Worker choices=SCRIPT_TEMPLATES.keys(), 239*8975f5c5SAndroid Build Coastguard Worker help='Language in which the wrapper script will be written.') 240*8975f5c5SAndroid Build Coastguard Worker parser.add_argument( 241*8975f5c5SAndroid Build Coastguard Worker 'executable_args', nargs='*', 242*8975f5c5SAndroid Build Coastguard Worker help='Arguments to wrap into the executable.') 243*8975f5c5SAndroid Build Coastguard Worker return parser 244*8975f5c5SAndroid Build Coastguard Worker 245*8975f5c5SAndroid Build Coastguard Worker 246*8975f5c5SAndroid Build Coastguard Workerdef main(raw_args): 247*8975f5c5SAndroid Build Coastguard Worker parser = CreateArgumentParser() 248*8975f5c5SAndroid Build Coastguard Worker args = parser.parse_args(raw_args) 249*8975f5c5SAndroid Build Coastguard Worker return Wrap(args) 250*8975f5c5SAndroid Build Coastguard Worker 251*8975f5c5SAndroid Build Coastguard Worker 252*8975f5c5SAndroid Build Coastguard Workerif __name__ == '__main__': 253*8975f5c5SAndroid Build Coastguard Worker sys.exit(main(sys.argv[1:])) 254