Accept a zip file on input

utils/kem-import.py

@@ -1,17 +1,19 @@
-# A Python script to decrypt the content of Kamstrup KEM file.
+# A Python script that decrypts the content of Kamstrup KEM file and imports meter files
+# to wmbusmeters' config folder.
 #
 # The KEM file is a (sometimes zipped) xml file that contains xml-encrypted data using 
 # the xml-enc standard (http://www.w3.org/2001/04/xmlenc). The password needed to decrypt 
 # the xml-encrypted data can either be the CustomerId or something else selected by the 
-# person that created the KEM file using Kamstrup software.
+# person that has created the KEM file using Kamstrup software.
 #
-# This script takes the encrypted KEM file and decrypts its content. The result is also 
-# a XML file with a list of meters. The script prints the information about each meter to
-# the console and populates wmbusmeters' config folder with corresponding meter files. 
+# This script takes the encrypted KEM file and decrypts its content (it automaticly detects
+# the zip archive and extracts the kem file from it). The result is a XML with a list of meters 
+# with their types, serial numbers, keys, etc. The script prints the information about each meter 
+# to the console and populates wmbusmeters' config folder with corresponding meter files. 
 # Optionally, decrypted KEM file content can be saved to a file (option -o).
 #
-# Usage: python kem-decryptor.py [options] <kem_file> <password>
-# kem_file ... the name of the KEM file to decrypt (you need to unzip the the KEM file first, if it is zipped)
+# Usage: python kem-import.py [options] <kem_file> <password>
+# kem_file ... the name of the KEM file to decrypt or a zip archive having the KEM file in it
 # password ... original password used to encrypt the content (16 characters maximum)
 # options  ... use -h switch to get a list of options with their descriptions
 #
@@ -26,15 +28,21 @@ import argparse
 import base64
 import Crypto.Cipher.AES as AES
 from xml.dom import minidom
+import zipfile
 
 help_prog = "Decrypts Kamstrup KEM file and imports meter information into wmbusmeters' config folder."
 help_epilog = ""
-help_kem = "The name of the KEM file to decrypt (you need to unzip the the KEM file first, if it is zipped)."
+help_kem = "The name of the KEM file to decrypt or a name of a zip archive that contains the encrypted KEM file."
 help_pwd = "The original password used to encrypt the KEM file content (16 characters maximum)."
-help_cfg = """Location of config files for wmbusmeters (default location is the current working directory). This option has the same meaning as --useconfig option of the wmbusmeters daemon, i.e. --useconfig=/ will populate /etc/wmbusmeters.d folder (must be run with sudo) and --useconfig=. (the default) populates ./etc/wmbusmeters.d folder. If the destination folder does not exist, it will be created; existing meter files will be overwritten."""
+help_cfg = """Location of config files for wmbusmeters (default location is the current working directory). 
+This option has the same meaning as --useconfig option of the wmbusmeters daemon, i.e. --useconfig=/ will 
+populate /etc/wmbusmeters.d folder (must be run with sudo) and --useconfig=. (the default) populates 
+./etc/wmbusmeters.d folder. If the destination folder does not exist, it will be created; existing meter 
+files will be overwritten."""
 help_dry = "No meter files will be created, only the info will be printed on the console."
 help_out = "Save the decrypted KEM file content into a given file."
 
+# define command line arguments
 argparser = argparse.ArgumentParser(description=help_prog, epilog=help_epilog)
 argparser.add_argument('kem_file', help=help_kem)
 argparser.add_argument('password', help=help_pwd)
@@ -42,17 +50,55 @@ argparser.add_argument("-c", "--useconfig", type=str, action='store', dest='conf
 argparser.add_argument("-n", "--dryrun", action='store_true', dest='dryrun', help=help_dry)
 argparser.add_argument("-o", "--output", type=str, action='store', dest='output', help=help_out)
 
+# parse command line arguments
 args = argparser.parse_args()
 
 # adjust the config folder location to full target path
 args.config = args.config.strip(os.path.sep) + '/etc/wmbusmeters.d/'
 
 
-# read and parse KEM file and extract its encrypted content
+# test if input file exists
+if (not os.path.isfile(args.kem_file)):
+    print('ERROR: The input KEM file does not exist.')
+    sys.exit(1)
+#end if
+
+kem_file_content = None
+
+# test if input file is a zip file
+#     yes: find a file with .kem extension in the zip and extract content of that file
+#     no : load the content of the input file directly
+if (zipfile.is_zipfile(args.kem_file)):
+    print("Detected a zip file on input ... extracting")
+    with zipfile.ZipFile(args.kem_file,'r') as zipobj:
+        file_list = zipobj.namelist()
+        for file_name in file_list:
+            if file_name.endswith('.kem'):
+                kem_file_content = zipobj.read(file_name)
+                break
+            #end if
+        #end for
+        
+        if (not kem_file_content):
+            print("ERROR: The zip file '%s' does not seem to contain any '.kem' file." % (args.kem_file))
+            sys.exit(1)
+        #end if
+else:
+    # read content of the kem file
+    with open(args.kem_file,'r') as f:
+        kem_file_content = f.read()
+#end if
+
+
+# read and parse KEM file content and extract its encrypted part
 # (KEM file is a a XML file formated according to http://www.w3.org/2001/04/xmlenc)
-xmldoc = minidom.parse(args.kem_file)
-encrypedtext = xmldoc.getElementsByTagName('CipherValue')[0].firstChild.nodeValue
-encrypeddata = base64.b64decode(encrypedtext)
+try:
+    xmldoc = minidom.parseString(kem_file_content)
+    encrypedtext = xmldoc.getElementsByTagName('CipherValue')[0].firstChild.nodeValue
+    encrypeddata = base64.b64decode(encrypedtext)
+except:
+    print('ERROR: The KEM file does not seem to contain encrypted data.')
+    sys.exit(1)
 
 # KEM file data is encrypted with AES-128-CBC cipher and decryption
 # requires a 128bit key and 128bit initialization vector. In the case of KEM file, 
@@ -138,7 +184,6 @@ for e in xmldoc.getElementsByTagName('Meter'):
                     print("You may need to use 'sudo' to access the target config folder.")
                     sys.exit(1)
                 else:
-                    print('XX',e)
                     raise
         else:
             print('    !! unknow meter type, meter file has not been created')