OLD | NEW |
1 # Copyright 2014 The LUCI Authors. All rights reserved. | 1 # Copyright 2014 The LUCI Authors. All rights reserved. |
2 # Use of this source code is governed under the Apache License, Version 2.0 | 2 # Use of this source code is governed under the Apache License, Version 2.0 |
3 # that can be found in the LICENSE file. | 3 # that can be found in the LICENSE file. |
4 | 4 |
5 """Generates the swarming_bot.zip archive for the bot. | 5 """Generates the swarming_bot.zip archive for the bot. |
6 | 6 |
7 Unlike the other source files, this file can be run from ../tools/bot_archive.py | 7 Unlike the other source files, this file can be run from ../tools/bot_archive.py |
8 stand-alone to generate a swarming_bot.zip for local testing so it doesn't | 8 stand-alone to generate a swarming_bot.zip for local testing so it doesn't |
9 import anything from the AppEngine SDK. | 9 import anything from the AppEngine SDK. |
10 | 10 |
(...skipping 359 matching lines...) Expand 10 before | Expand all | Expand 10 after Loading... |
370 partial = os.path.sep.join(parts[:i]) | 370 partial = os.path.sep.join(parts[:i]) |
371 if os.path.isfile(partial): | 371 if os.path.isfile(partial): |
372 with open(partial) as f: | 372 with open(partial) as f: |
373 link = f.read() | 373 link = f.read() |
374 assert '\n' not in link and link, link | 374 assert '\n' not in link and link, link |
375 parts[i-1] = link | 375 parts[i-1] = link |
376 return os.path.normpath(os.path.sep.join(parts)) | 376 return os.path.normpath(os.path.sep.join(parts)) |
377 | 377 |
378 | 378 |
379 def yield_swarming_bot_files( | 379 def yield_swarming_bot_files( |
380 root_dir, host, host_version, additionals, enable_ts_monitoring): | 380 root_dir, host, host_version, additionals, settings): |
381 """Yields all the files to map as tuple(filename, content). | 381 """Yields all the files to map as tuple(filename, content). |
382 | 382 |
383 config.json is injected with json data about the server. | 383 config.json is injected with json data about the server. |
384 | 384 |
385 This function guarantees that the output is sorted by filename. | 385 This function guarantees that the output is sorted by filename. |
386 """ | 386 """ |
387 grpc_prefix = 'grpc://' | |
388 is_grpc = host.startswith(grpc_prefix) | |
389 if is_grpc: | |
390 host = host[len(grpc_prefix):] | |
391 | |
392 items = {i: None for i in FILES} | 387 items = {i: None for i in FILES} |
393 items.update(additionals) | 388 items.update(additionals) |
394 config = { | 389 config = { |
395 'enable_ts_monitoring': enable_ts_monitoring, | 390 'enable_ts_monitoring': settings.enable_ts_monitoring, |
396 'is_grpc': is_grpc, | |
397 'server': host.rstrip('/'), | 391 'server': host.rstrip('/'), |
398 'server_version': host_version, | 392 'server_version': host_version, |
| 393 'isolate_grpc_proxy': settings.bot_isolate_grpc_proxy, |
399 } | 394 } |
400 items['config/config.json'] = json.dumps(config) | 395 items['config/config.json'] = json.dumps(config) |
| 396 logging.debug('Bot config.json: %s', items['config/config.json']) |
401 for item, content in sorted(items.iteritems()): | 397 for item, content in sorted(items.iteritems()): |
402 if content is not None: | 398 if content is not None: |
403 yield item, content | 399 yield item, content |
404 else: | 400 else: |
405 with open(resolve_symlink(os.path.join(root_dir, item)), 'rb') as f: | 401 with open(resolve_symlink(os.path.join(root_dir, item)), 'rb') as f: |
406 yield item, f.read() | 402 yield item, f.read() |
407 | 403 |
408 | 404 |
409 def get_swarming_bot_zip( | 405 def get_swarming_bot_zip( |
410 root_dir, host, host_version, additionals, enable_ts_monitoring): | 406 root_dir, host, host_version, additionals, settings): |
411 """Returns a zipped file of all the files a bot needs to run. | 407 """Returns a zipped file of all the files a bot needs to run. |
412 | 408 |
413 Arguments: | 409 Arguments: |
414 root_dir: directory swarming_bot. | 410 root_dir: directory swarming_bot. |
415 additionals: dict(filepath: content) of additional items to put into the zip | 411 additionals: dict(filepath: content) of additional items to put into the zip |
416 file, in addition to FILES and MAPPED. In practice, it's going to be a | 412 file, in addition to FILES and MAPPED. In practice, it's going to be a |
417 custom bot_config.py. | 413 custom bot_config.py. |
418 enable_ts_monitoring: bool if ts_mon should be enabled on the bot. | 414 enable_ts_monitoring: bool if ts_mon should be enabled on the bot. |
419 | 415 |
420 Returns: | 416 Returns: |
421 Tuple(str being the zipped file's content, bot version (SHA256) it | 417 Tuple(str being the zipped file's content, bot version (SHA256) it |
422 represents). | 418 represents). |
423 """ | 419 """ |
424 zip_memory_file = StringIO.StringIO() | 420 zip_memory_file = StringIO.StringIO() |
425 h = hashlib.sha256() | 421 h = hashlib.sha256() |
426 with zipfile.ZipFile(zip_memory_file, 'w', zipfile.ZIP_DEFLATED) as zip_file: | 422 with zipfile.ZipFile(zip_memory_file, 'w', zipfile.ZIP_DEFLATED) as zip_file: |
427 for name, content in yield_swarming_bot_files( | 423 for name, content in yield_swarming_bot_files( |
428 root_dir, host, host_version, additionals, enable_ts_monitoring): | 424 root_dir, host, host_version, additionals, settings): |
429 zip_file.writestr(name, content) | 425 zip_file.writestr(name, content) |
430 h.update(str(len(name))) | 426 h.update(str(len(name))) |
431 h.update(name) | 427 h.update(name) |
432 h.update(str(len(content))) | 428 h.update(str(len(content))) |
433 h.update(content) | 429 h.update(content) |
434 | 430 |
435 data = zip_memory_file.getvalue() | 431 data = zip_memory_file.getvalue() |
436 bot_version = h.hexdigest() | 432 bot_version = h.hexdigest() |
437 logging.info( | 433 logging.info( |
438 'get_swarming_bot_zip(%s) is %d bytes; %s', | 434 'get_swarming_bot_zip(%s) is %d bytes; %s', |
439 additionals.keys(), len(data), bot_version) | 435 additionals.keys(), len(data), bot_version) |
440 return data, bot_version | 436 return data, bot_version |
441 | 437 |
442 | 438 |
443 def get_swarming_bot_version( | 439 def get_swarming_bot_version( |
444 root_dir, host, host_version, additionals, enable_ts_monitoring): | 440 root_dir, host, host_version, additionals, settings): |
445 """Returns the SHA256 hash of the bot code, representing the version. | 441 """Returns the SHA256 hash of the bot code, representing the version. |
446 | 442 |
447 Arguments: | 443 Arguments: |
448 root_dir: directory swarming_bot. | 444 root_dir: directory swarming_bot. |
449 additionals: See get_swarming_bot_zip's doc. | 445 additionals: See get_swarming_bot_zip's doc. |
450 enable_ts_monitoring: bool if ts_mon should be enabled on the bot. | 446 enable_ts_monitoring: bool if ts_mon should be enabled on the bot. |
451 | 447 |
452 Returns: | 448 Returns: |
453 The SHA256 hash of the bot code. | 449 The SHA256 hash of the bot code. |
454 """ | 450 """ |
455 h = hashlib.sha256() | 451 h = hashlib.sha256() |
456 try: | 452 try: |
457 # TODO(maruel): Deduplicate from zip_package.genereate_version(). | 453 # TODO(maruel): Deduplicate from zip_package.genereate_version(). |
458 for name, content in yield_swarming_bot_files( | 454 for name, content in yield_swarming_bot_files( |
459 root_dir, host, host_version, additionals, enable_ts_monitoring): | 455 root_dir, host, host_version, additionals, settings): |
460 h.update(str(len(name))) | 456 h.update(str(len(name))) |
461 h.update(name) | 457 h.update(name) |
462 h.update(str(len(content))) | 458 h.update(str(len(content))) |
463 h.update(content) | 459 h.update(content) |
464 except IOError: | 460 except IOError: |
465 logging.warning('Missing expected file. Hash will be invalid.') | 461 logging.warning('Missing expected file. Hash will be invalid.') |
466 bot_version = h.hexdigest() | 462 bot_version = h.hexdigest() |
467 logging.info( | 463 logging.info( |
468 'get_swarming_bot_version(%s) = %s', sorted(additionals), bot_version) | 464 'get_swarming_bot_version(%s) = %s', sorted(additionals), bot_version) |
469 return bot_version | 465 return bot_version |
OLD | NEW |