| 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 |