########################################################################
# Copyright (c) 2010-2024 Broadcom. All Rights Reserved.
# Broadcom Confidential. The term "Broadcom" refers to Broadcom Inc.
# and/or its subsidiaries.
########################################################################

"""This module provides the gzip interface for ESXi"""

import gzip

# Static header for ESX gzip format: 2-byte gzip magic + compression
# method (8, "deflate") + flags (0) + 4-byte timestamp (0's) + extra
# flags (2, max compression) + OS (255, unknown). All optional fields are
# intentionally omitted. This enables re-compressing data and getting
# consistent output that can be matched against a digest.
ESX_GZIP_HEADER = b"\037\213\010\000\000\000\000\000\002\377"

class GzipFile(gzip.GzipFile):
   def _write_gzip_header(self, compresslevel=9):
      self.fileobj.write(ESX_GZIP_HEADER)

def main():
   """Entrypoint for ESX payload builds.
   """
   # NOTE: Keep "build-specific" imports here.
   import argparse
   import shutil

   parser = argparse.ArgumentParser(
      description='GZIP-compress ESXi VIB payload file',
   )
   parser.add_argument(
      '--input', help="Input file [stdin]",
      type=argparse.FileType('rb'), required=True, default='-',
   )
   parser.add_argument(
      '--compresslevel', help="Compress level of gzip [stdin]",
      type=int, default=9,
   )
   parser.add_argument(
      '--output', help="Output file [stdout]",
      type=argparse.FileType('wb'), required=True, default='-',
   )

   args = parser.parse_args()

   with GzipFile(mode='wb', compresslevel=args.compresslevel,
                 fileobj=args.output) as output:
      shutil.copyfileobj(args.input, output)

if __name__ == '__main__':
   main()

