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