Fix lint and typecheck warnings

sevsnpmeasure/cli.py

@@ -19,6 +19,18 @@ def auto_base_int(s: str) -> int:
     return int(s, 0)
 
 
+def print_measurement(ld: bytes, sev_mode: SevMode, output_format: str, verbose: bool):
+    if output_format == "hex":
+        measurement = ld.hex()
+    elif output_format == "base64":
+        measurement = base64.b64encode(ld).decode()
+
+    if verbose:
+        print(f"Calculated {sev_mode.name} guest measurement: {measurement}")
+    else:
+        print(measurement)
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(prog='sev-snp-measure',
                                      description='Calculate AMD SEV/SEV-ES/SEV-SNP guest launch measurement')
@@ -34,7 +46,7 @@ def main() -> int:
     parser.add_argument('--vcpu-model', metavar='MODEL', type=int, help='Guest vcpu model', default=None)
     parser.add_argument('--vcpu-stepping', metavar='STEPPING', type=int, help='Guest vcpu stepping', default=None)
     parser.add_argument('--vmm-type', metavar='VMMTYPE', type=str,
-                        help=f"Type of guest vmm ({', '.join(vmm_types.VMMType._member_names_)})", default='QEMU')
+                        help=f"Type of guest vmm ({', '.join(vmm_types.VMMType.__members__.keys())})", default='QEMU')
     parser.add_argument('--ovmf', metavar='PATH',
                         help='OVMF file to calculate hash from', required=True)
     parser.add_argument('--kernel', metavar='PATH',
@@ -54,7 +66,7 @@ def main() -> int:
     if args.mode != 'sev' and args.vcpus is None:
         parser.error(f"missing --vcpus N in guest mode '{args.mode}'")
 
-    if args.vmm_type in vmm_types.VMMType._member_names_:
+    if args.vmm_type in vmm_types.VMMType.__members__.keys():
         vmm_type = vmm_types.VMMType[args.vmm_type]
     else:
         parser.error(f"unknown VMM type '{args.vmm_type}'")
@@ -73,17 +85,9 @@ def main() -> int:
     sev_mode = SevMode.from_str(args.mode)
     ld = guest.calc_launch_digest(sev_mode, args.vcpus, vcpu_sig, args.ovmf,
                                   args.kernel, args.initrd, args.append, args.snp_ovmf_hash,
-                                  vmm_type = vmm_type)
-
-    if args.output_format == "hex":
-        measurement = ld.hex()
-    elif args.output_format == "base64":
-        measurement = base64.b64encode(ld).decode()
+                                  vmm_type=vmm_type)
 
-    if args.verbose:
-        print(f"Calculated {sev_mode.name} guest measurement: {measurement}")
-    else:
-        print(measurement)
+    print_measurement(ld, sev_mode, args.output_format, args.verbose)
     return 0
 
 

sevsnpmeasure/guest.py

@@ -25,7 +25,7 @@ def calc_launch_digest(mode: SevMode, vcpus: int, vcpu_sig: int, ovmf_file: str,
     if mode == SevMode.SEV_SNP:
         return snp_calc_launch_digest(vcpus, vcpu_sig, ovmf_file, kernel, initrd, append, snp_ovmf_hash_str, vmm_type)
     elif mode == SevMode.SEV_ES:
-        return seves_calc_launch_digest(vcpus, vcpu_sig, ovmf_file, kernel, initrd, append, vmm_type = vmm_type)
+        return seves_calc_launch_digest(vcpus, vcpu_sig, ovmf_file, kernel, initrd, append, vmm_type=vmm_type)
     elif mode == SevMode.SEV:
         return sev_calc_launch_digest(ovmf_file, kernel, initrd, append)
     else:
@@ -43,7 +43,7 @@ def snp_update_kernel_hashes(gctx: GCTX, ovmf: OVMF, sev_hashes: Optional[SevHas
         gctx.update_zero_pages(gpa, size)
 
 
-def snp_update_metadata_pages(gctx: GCTX, ovmf: OVMF, sev_hashes: Optional[SevHashes], vmm_type: VMMType = VMMType.QEMU) -> None:
+def snp_update_metadata_pages(gctx: GCTX, ovmf: OVMF, sev_hashes: Optional[SevHashes], vmm_type: VMMType) -> None:
     for desc in ovmf.metadata_items():
         if desc.section_type() == SectionType.SNP_SEC_MEM:
             gctx.update_zero_pages(desc.gpa, desc.size)

sevsnpmeasure/vmm_types.py

@@ -1,5 +1,6 @@
 from enum import Enum
 
+
 class VMMType(Enum):
     QEMU = 1
     ec2 = 2

sevsnpmeasure/vmsa.py

@@ -7,7 +7,6 @@
 from ctypes import c_uint8, c_uint16, c_uint32, c_uint64
 from typing import Iterator
 from .sev_mode import SevMode
-from sevsnpmeasure import vcpu_types
 from .vmm_types import VMMType
 
 

tests/test_guest.py

@@ -61,7 +61,7 @@ def test_snp_ec2(self):
                 "/dev/null",
                 "/dev/null",
                 "",
-                vmm_type = vmm_types.VMMType.ec2)
+                vmm_type=vmm_types.VMMType.ec2)
         self.assertEqual(
                 ld.hex(),
                 '21408991204a77c12eb014cc7a1ad723b423069f421b873d'