#!/usr/bin/env python

"""
Copyright 2020 VMware, Inc.  All rights reserved. -- VMware Confidential

Script to fix the configSchemas inside metadata.zip in patch-store.
"""

import os
import pwd
import sys
import logging
import tempfile
import zipfile

sys.path.append('/usr/lib/vmware-imagebuilder/site-packages')

from vmware.esximage.ConfigSchema import ConfigSchema

def fixMetadataZips(src_dir='/storage/updatemgr/patch-store'):
   if not os.path.exists(src_dir):
      return

   pwnam = pwd.getpwnam('updatemgr')
   uid = pwnam.pw_uid
   gid = pwnam.pw_gid
   for root, dirs, files in os.walk(src_dir):
      for f in files:
         oldfile = os.path.join(root, f)
         if not (os.path.isfile(oldfile) and oldfile.endswith('.zip') and
             'metadata' in oldfile):
            continue

         newfile = ''
         try:
            fd, newfile = tempfile.mkstemp(dir=root)
            os.close(fd)
            schemas = {}
            with zipfile.ZipFile(oldfile, 'r') as zipin:
               for i in zipin.infolist():
                  if not (i.filename.startswith('configSchemas/') and
                          i.filename.endswith('json')):
                     continue

                  try:
                     cs = zipin.read(i.filename).decode('utf-8')
                     filteredCs = ConfigSchema(cs)._filteredSchemaStr

                     if not filteredCs or filteredCs == cs:
                        # There is nothing to filter
                        continue
                  except:
                     logging.exception("Failed to read configSchema: %s",
                                       i.filename)
                     continue

                  logging.info("Filtering %s(%s)", oldfile, i.filename)
                  schemas[i.filename] = filteredCs

               if not schemas:
                  continue

               with zipfile.ZipFile(newfile, 'w') as zipout:
                  zipout.comment = zipin.comment
                  for i in zipin.infolist():
                     if i.filename not in schemas:
                        zipout.writestr(i, zipin.read(i.filename))
                     else:
                        zipout.writestr(i, schemas[i.filename])

            os.replace(newfile, oldfile)
            os.chown(oldfile, uid, gid)
         except:
            logging.exception("Failed to fix the metadata in file: %s",
                              oldfile)
         finally:
            try:
               os.remove(newfile)
            except FileNotFoundError:
               pass
