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