my_init
1 #!/usr/bin/python3 -u 2 # -*- coding: utf-8 -*- 3 4 import argparse 5 import errno 6 import json 7 import os 8 import os.path 9 import re 10 import signal 11 import stat 12 import sys 13 import time 14 15 ENV_INIT_DIRECTORY = os.environ.get('ENV_INIT_DIRECTORY', '/etc/my_init.d') 16 17 KILL_PROCESS_TIMEOUT = int(os.environ.get('KILL_PROCESS_TIMEOUT', 30)) 18 KILL_ALL_PROCESSES_TIMEOUT = int(os.environ.get('KILL_ALL_PROCESSES_TIMEOUT', 30)) 19 20 LOG_LEVEL_ERROR = 1 21 LOG_LEVEL_WARN = 1 22 LOG_LEVEL_INFO = 2 23 LOG_LEVEL_DEBUG = 3 24 25 SHENV_NAME_WHITELIST_REGEX = re.compile('\W') 26 27 log_level = None 28 29 terminated_child_processes = {} 30 31 _find_unsafe = re.compile(r'[^\w@%+=:,./-]').search 32 33 34 class AlarmException(Exception): 35 pass 36 37 38 def error(message): 39 if log_level >= LOG_LEVEL_ERROR: 40 sys.stderr.write("*** %s\n" % message) 41 42 43 def warn(message): 44 if log_level >= LOG_LEVEL_WARN: 45 sys.stderr.write("*** %s\n" % message) 46 47 48 def info(message): 49 if log_level >= LOG_LEVEL_INFO: 50 sys.stderr.write("*** %s\n" % message) 51 52 53 def debug(message): 54 if log_level >= LOG_LEVEL_DEBUG: 55 sys.stderr.write("*** %s\n" % message) 56 57 58 def ignore_signals_and_raise_keyboard_interrupt(signame): 59 signal.signal(signal.SIGTERM, signal.SIG_IGN) 60 signal.signal(signal.SIGINT, signal.SIG_IGN) 61 raise KeyboardInterrupt(signame) 62 63 64 def raise_alarm_exception(): 65 raise AlarmException('Alarm') 66 67 68 def listdir(path): 69 try: 70 result = os.stat(path) 71 except OSError: 72 return [] 73 if stat.S_ISDIR(result.st_mode): 74 return sorted(os.listdir(path)) 75 else: 76 return [] 77 78 79 def is_exe(path): 80 try: 81 return os.path.isfile(path) and os.access(path, os.X_OK) 82 except OSError: 83 return False 84 85 86 def import_envvars(clear_existing_environment=True, override_existing_environment=True): 87 if not os.path.exists("/etc/container_environment"): 88 return 89 new_env = {} 90 for envfile in listdir("/etc/container_environment"): 91 name = os.path.basename(envfile) 92 with open("/etc/container_environment/" + envfile, "r") as f: 93 # Text files often end with a trailing newline, which we 94 # don't want to include in the env variable value. See 95 # https://github.com/phusion/baseimage-docker/pull/49 96 value = re.sub('\n\Z', '', f.read()) 97 new_env[name] = value 98 if clear_existing_environment: 99 os.environ.clear() 100 for name, value in new_env.items(): 101 if override_existing_environment or name not in os.environ: 102 os.environ[name] = value 103 104 105 def export_envvars(to_dir=True): 106 if not os.path.exists("/etc/container_environment"): 107 return 108 shell_dump = "" 109 for name, value in os.environ.items(): 110 if name in ['HOME', 'USER', 'GROUP', 'UID', 'GID', 'SHELL']: 111 continue 112 if to_dir: 113 with open("/etc/container_environment/" + name, "w") as f: 114 f.write(value) 115 shell_dump += "export " + sanitize_shenvname(name) + "=" + shquote(value) + "\n" 116 with open("/etc/container_environment.sh", "w") as f: 117 f.write(shell_dump) 118 with open("/etc/container_environment.json", "w") as f: 119 f.write(json.dumps(dict(os.environ))) 120 121 122 def shquote(s): 123 """Return a shell-escaped version of the string *s*.""" 124 if not s: 125 return "''" 126 if _find_unsafe(s) is None: 127 return s 128 129 # use single quotes, and put single quotes into double quotes 130 # the string $'b is then quoted as '$'"'"'b' 131 return "'" + s.replace("'", "'\"'\"'") + "'" 132 133 134 def sanitize_shenvname(s): 135 """Return string with [0-9a-zA-Z_] characters""" 136 return re.sub(SHENV_NAME_WHITELIST_REGEX, "_", s) 137 138 139 # Waits for the child process with the given PID, while at the same time 140 # reaping any other child processes that have exited (e.g. adopted child 141 # processes that have terminated). 142 143 def waitpid_reap_other_children(pid): 144 global terminated_child_processes 145 146 status = terminated_child_processes.get(pid) 147 if status: 148 # A previous call to waitpid_reap_other_children(), 149 # with an argument not equal to the current argument, 150 # already waited for this process. Return the status 151 # that was obtained back then. 152 del terminated_child_processes[pid] 153 return status 154 155 done = False 156 status = None 157 while not done: 158 try: 159 # https://github.com/phusion/baseimage-docker/issues/151#issuecomment-92660569 160 this_pid, status = os.waitpid(pid, os.WNOHANG) 161 if this_pid == 0: 162 this_pid, status = os.waitpid(-1, 0) 163 if this_pid == pid: 164 done = True 165 else: 166 # Save status for later. 167 terminated_child_processes[this_pid] = status 168 except OSError as e: 169 if e.errno == errno.ECHILD or e.errno == errno.ESRCH: 170 return None 171 else: 172 raise 173 return status 174 175 176 def stop_child_process(name, pid, signo=signal.SIGTERM, time_limit=KILL_PROCESS_TIMEOUT): 177 info("Shutting down %s (PID %d)..." % (name, pid)) 178 try: 179 os.kill(pid, signo) 180 except OSError: 181 pass 182 signal.alarm(time_limit) 183 try: 184 try: 185 waitpid_reap_other_children(pid) 186 except OSError: 187 pass 188 except AlarmException: 189 warn("%s (PID %d) did not shut down in time. Forcing it to exit." % (name, pid)) 190 try: 191 os.kill(pid, signal.SIGKILL) 192 except OSError: 193 pass 194 try: 195 waitpid_reap_other_children(pid) 196 except OSError: 197 pass 198 finally: 199 signal.alarm(0) 200 201 202 def run_command_killable(*argv): 203 filename = argv[0] 204 status = None 205 pid = os.spawnvp(os.P_NOWAIT, filename, argv) 206 try: 207 status = waitpid_reap_other_children(pid) 208 except BaseException: 209 warn("An error occurred. Aborting.") 210 stop_child_process(filename, pid) 211 raise 212 if status != 0: 213 if status is None: 214 error("%s exited with unknown status\n" % filename) 215 else: 216 error("%s failed with status %d\n" % (filename, os.WEXITSTATUS(status))) 217 sys.exit(1) 218 219 220 def run_command_killable_and_import_envvars(*argv): 221 run_command_killable(*argv) 222 import_envvars() 223 export_envvars(False) 224 225 226 def kill_all_processes(time_limit): 227 info("Killing all processes...") 228 try: 229 os.kill(-1, signal.SIGTERM) 230 except OSError: 231 pass 232 signal.alarm(time_limit) 233 try: 234 # Wait until no more child processes exist. 235 done = False 236 while not done: 237 try: 238 os.waitpid(-1, 0) 239 except OSError as e: 240 if e.errno == errno.ECHILD: 241 done = True 242 else: 243 raise 244 except AlarmException: 245 warn("Not all processes have exited in time. Forcing them to exit.") 246 try: 247 os.kill(-1, signal.SIGKILL) 248 except OSError: 249 pass 250 finally: 251 signal.alarm(0) 252 253 254 def run_startup_files(): 255 # Run ENV_INIT_DIRECTORY/* 256 for name in listdir(ENV_INIT_DIRECTORY): 257 filename = os.path.join(ENV_INIT_DIRECTORY, name) 258 if is_exe(filename): 259 info("Running %s..." % filename) 260 run_command_killable_and_import_envvars(filename) 261 262 # Run /etc/rc.local. 263 if is_exe("/etc/rc.local"): 264 info("Running /etc/rc.local...") 265 run_command_killable_and_import_envvars("/etc/rc.local") 266 267 268 def run_pre_shutdown_scripts(): 269 debug("Running pre-shutdown scripts...") 270 271 # Run /etc/my_init.pre_shutdown.d/* 272 for name in listdir("/etc/my_init.pre_shutdown.d"): 273 filename = "/etc/my_init.pre_shutdown.d/" + name 274 if is_exe(filename): 275 info("Running %s..." % filename) 276 run_command_killable(filename) 277 278 279 def run_post_shutdown_scripts(): 280 debug("Running post-shutdown scripts...") 281 282 # Run /etc/my_init.post_shutdown.d/* 283 for name in listdir("/etc/my_init.post_shutdown.d"): 284 filename = "/etc/my_init.post_shutdown.d/" + name 285 if is_exe(filename): 286 info("Running %s..." % filename) 287 run_command_killable(filename) 288 289 290 def start_runit(): 291 info("Booting runit daemon...") 292 pid = os.spawnl(os.P_NOWAIT, "/usr/bin/runsvdir", "/usr/bin/runsvdir", 293 "-P", "/etc/service") 294 info("Runit started as PID %d" % pid) 295 return pid 296 297 298 def wait_for_runit_or_interrupt(pid): 299 status = waitpid_reap_other_children(pid) 300 return (True, status) 301 302 303 def shutdown_runit_services(quiet=False): 304 if not quiet: 305 debug("Begin shutting down runit services...") 306 os.system("/usr/bin/sv -w %d force-stop /etc/service/* > /dev/null" % KILL_PROCESS_TIMEOUT) 307 308 309 def wait_for_runit_services(): 310 debug("Waiting for runit services to exit...") 311 done = False 312 while not done: 313 done = os.system("/usr/bin/sv status /etc/service/* | grep -q '^run:'") != 0 314 if not done: 315 time.sleep(0.1) 316 # According to https://github.com/phusion/baseimage-docker/issues/315 317 # there is a bug or race condition in Runit, causing it 318 # not to shutdown services that are already being started. 319 # So during shutdown we repeatedly instruct Runit to shutdown 320 # services. 321 shutdown_runit_services(True) 322 323 324 def install_insecure_key(): 325 info("Installing insecure SSH key for user root") 326 run_command_killable("/usr/sbin/enable_insecure_key") 327 328 329 def main(args): 330 import_envvars(False, False) 331 export_envvars() 332 333 if args.enable_insecure_key: 334 install_insecure_key() 335 336 if not args.skip_startup_files: 337 run_startup_files() 338 339 runit_exited = False 340 exit_code = None 341 342 if not args.skip_runit: 343 runit_pid = start_runit() 344 try: 345 exit_status = None 346 if len(args.main_command) == 0: 347 runit_exited, exit_code = wait_for_runit_or_interrupt(runit_pid) 348 if runit_exited: 349 if exit_code is None: 350 info("Runit exited with unknown status") 351 exit_status = 1 352 else: 353 exit_status = os.WEXITSTATUS(exit_code) 354 info("Runit exited with status %d" % exit_status) 355 else: 356 info("Running %s..." % " ".join(args.main_command)) 357 pid = os.spawnvp(os.P_NOWAIT, args.main_command[0], args.main_command) 358 try: 359 exit_code = waitpid_reap_other_children(pid) 360 if exit_code is None: 361 info("%s exited with unknown status." % args.main_command[0]) 362 exit_status = 1 363 else: 364 exit_status = os.WEXITSTATUS(exit_code) 365 info("%s exited with status %d." % (args.main_command[0], exit_status)) 366 except KeyboardInterrupt: 367 stop_child_process(args.main_command[0], pid) 368 raise 369 except BaseException: 370 warn("An error occurred. Aborting.") 371 stop_child_process(args.main_command[0], pid) 372 raise 373 sys.exit(exit_status) 374 finally: 375 if not args.skip_runit: 376 run_pre_shutdown_scripts() 377 shutdown_runit_services() 378 if not runit_exited: 379 stop_child_process("runit daemon", runit_pid) 380 wait_for_runit_services() 381 run_post_shutdown_scripts() 382 383 # Parse options. 384 parser = argparse.ArgumentParser(description='Initialize the system.') 385 parser.add_argument('main_command', metavar='MAIN_COMMAND', type=str, nargs='*', 386 help='The main command to run. (default: runit)') 387 parser.add_argument('--enable-insecure-key', dest='enable_insecure_key', 388 action='store_const', const=True, default=False, 389 help='Install the insecure SSH key') 390 parser.add_argument('--skip-startup-files', dest='skip_startup_files', 391 action='store_const', const=True, default=False, 392 help='Skip running /etc/my_init.d/* and /etc/rc.local') 393 parser.add_argument('--skip-runit', dest='skip_runit', 394 action='store_const', const=True, default=False, 395 help='Do not run runit services') 396 parser.add_argument('--no-kill-all-on-exit', dest='kill_all_on_exit', 397 action='store_const', const=False, default=True, 398 help='Don\'t kill all processes on the system upon exiting') 399 parser.add_argument('--quiet', dest='log_level', 400 action='store_const', const=LOG_LEVEL_WARN, default=LOG_LEVEL_INFO, 401 help='Only print warnings and errors') 402 args = parser.parse_args() 403 log_level = args.log_level 404 405 if args.skip_runit and len(args.main_command) == 0: 406 error("When --skip-runit is given, you must also pass a main command.") 407 sys.exit(1) 408 409 # Run main function. 410 signal.signal(signal.SIGTERM, lambda signum, frame: ignore_signals_and_raise_keyboard_interrupt('SIGTERM')) 411 signal.signal(signal.SIGINT, lambda signum, frame: ignore_signals_and_raise_keyboard_interrupt('SIGINT')) 412 signal.signal(signal.SIGALRM, lambda signum, frame: raise_alarm_exception()) 413 try: 414 main(args) 415 except KeyboardInterrupt: 416 warn("Init system aborted.") 417 exit(2) 418 finally: 419 if args.kill_all_on_exit: 420 kill_all_processes(KILL_ALL_PROCESSES_TIMEOUT)