Merge pull request #10 from NBISweden/header

Full header = (start, enc part)

legacryptor/crypt4gh.py

@@ -210,16 +210,18 @@ def get_header(infile):
     '''Extract header and advance file position to AES block.'''
 
     LOG.info(f'Deconstructing the Header')
-    magic_number = infile.read(8)
+    buf = bytearray(16)
+    infile.readinto(buf)
+    magic_number = buf[:8]
     if magic_number != MAGIC_NUMBER:
         raise InvalidFormatError()
 
-    version = int.from_bytes(infile.read(4), byteorder='little')
+    version = int.from_bytes(buf[8:12], byteorder='little')
     if version != __version__:
         raise VersionError(version)
 
-    length = int.from_bytes(infile.read(4), byteorder='little') - 16
-    return infile.read(length)
+    length = int.from_bytes(buf[12:16], byteorder='little') - 16
+    return (bytes(buf), infile.read(length))
 
 # That allows us to decrypt and:
 # - dump the output to a file
@@ -271,7 +273,7 @@ def decrypt(privkey, infile, process_output=do_nothing, chunk_size=4096):
     assert privkey.is_unlocked, "The private key should be unlocked"
     assert chunk_size >= 32, "Chunk size larger than 32 bytes required"
 
-    encrypted_part = get_header(infile)
+    _, encrypted_part = get_header(infile)
     header = Header.decrypt(encrypted_part, privkey)
     # Only interested in the first record, for the moment
     r = header.records[0]
@@ -285,7 +287,7 @@ def reencrypt(pubkey, privkey, infile, process_output=do_nothing, chunk_size=409
     assert privkey.is_unlocked, "The private key should be unlocked"
     assert chunk_size >= 32, "Chunk size larger than 32 bytes required"
 
-    encrypted_part = get_header(infile)
+    _, encrypted_part = get_header(infile)
     header = Header.decrypt(encrypted_part, privkey)
     header_bytes = header.encrypt(pubkey)
     process_output(header_bytes)

tests/test_crypt4gh.py

@@ -84,7 +84,7 @@ def test_header(self, filedir):
         sec_keyfile = filedir.write('sec_key.asc', pgp_data.PGP_PRIVKEY.encode('utf-8'))
         sec_key, _ = pgpy.PGPKey.from_file(sec_keyfile)
         with sec_key.unlock(pgp_data.PGP_PASSPHRASE) as privkey:
-            header = Header.decrypt(get_header(open(infile, 'rb')), privkey)
+            header = Header.decrypt(get_header(open(infile, 'rb'))[1], privkey)
             self.assertEqual(pgp_data.RECORD_HEADER, str(header.records[0]))
             self.assertEqual(pgp_data.SESSION_KEY, header.records[0].session_key.hex())
             self.assertEqual(pgp_data.RECORD_HEADER_REPR, repr(header))
@@ -101,15 +101,15 @@ def test_header_bad_version(self, filedir):
     def test_key_id(self, filedir):
         """Testing get_key_id, should return the KeyID."""
         infile = filedir.write('infile.in', bytearray.fromhex(pgp_data.ENC_FILE))
-        header = get_header(open(infile, 'rb'))
+        _, header = get_header(open(infile, 'rb'))
         self.assertEqual(pgp_data.KEY_ID, get_key_id(header))
         filedir.cleanup()
 
     @tempdir()
     def test_header_to_records(self, filedir):
         """Should return one header record."""
         infile = filedir.write('infile.in', bytearray.fromhex(pgp_data.ENC_FILE))
-        header = get_header(open(infile, 'rb'))
+        _, header = get_header(open(infile, 'rb'))
         result = header_to_records(pgp_data.PGP_PRIVKEY, header, pgp_data.PGP_PASSPHRASE)
         self.assertEqual(1, len(result))
         self.assertEqual(pgp_data.RECORD_HEADER, str(result[0]))