From 5e75fffa864b323c2553a2ba1c3e0ae90a25597b Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Sun, 5 May 2019 15:41:47 -0600
Subject: [PATCH] Added 4K alignment check to checklists

---
 .bin/Scripts/functions/sw_diags.py  | 33 +++++++++++++++++++++++++++++
 .bin/Scripts/new_system_setup.py    |  4 ++++
 .bin/Scripts/system_checklist.py    | 12 +++++++----
 .bin/Scripts/system_checklist_hw.py | 12 +++++++----
 4 files changed, 53 insertions(+), 8 deletions(-)

diff --git a/.bin/Scripts/functions/sw_diags.py b/.bin/Scripts/functions/sw_diags.py
index 3aa2838c..0e73195e 100644
--- a/.bin/Scripts/functions/sw_diags.py
+++ b/.bin/Scripts/functions/sw_diags.py
@@ -6,6 +6,39 @@ from functions.common import *
 from settings.sw_diags import *
 
 
+class Not4KAlignedError(Exception):
+  pass
+
+
+def check_4k_alignment(show_alert=False):
+  """Check that all partitions are 4K aligned."""
+  aligned = True
+  cmd = ['WMIC', 'partition', 'get', 'StartingOffset']
+  offsets = []
+
+  # Get offsets
+  result = run_program(cmd, encoding='utf-8', errors='ignore', check=False)
+  offsets = result.stdout.splitlines()
+
+  # Check offsets
+  for off in offsets:
+    off = off.strip()
+    if not off.isnumeric():
+      # Skip
+      continue
+
+    try:
+      aligned &= int(off) % 4096 == 0
+    except ValueError:
+      # Ignore, this check is low priority
+      pass
+
+  # Show alert
+  if show_alert:
+    show_alert_box('One or more partitions are not 4K aligned')
+    raise Not4KAlignedError
+
+
 def check_connection():
   """Check if the system is online and optionally abort the script."""
   while True:
diff --git a/.bin/Scripts/new_system_setup.py b/.bin/Scripts/new_system_setup.py
index b0e9dbcd..2ed51b8b 100644
--- a/.bin/Scripts/new_system_setup.py
+++ b/.bin/Scripts/new_system_setup.py
@@ -23,6 +23,7 @@ if __name__ == '__main__':
       'CalledProcessError':       'Unknown Error',
       'FileNotFoundError':        'File not found',
       'GenericError':             'Unknown Error',
+      'Not4KAlignedError':        'False',
       'SecureBootDisabledError':  'Disabled',
     },
     'Warning': {
@@ -172,6 +173,9 @@ if __name__ == '__main__':
     try_and_print(message='Installed Antivirus:',
       function=get_installed_antivirus, ns='Unknown',
       other_results=other_results, print_return=True)
+    try_and_print(message='Partitions 4K aligned:',
+      function=check_4k_alignment, cs='True',
+      other_results=other_results)
 
     # Play audio, show devices, open Windows updates, and open Activation
     try_and_print(message='Opening Device Manager...',
diff --git a/.bin/Scripts/system_checklist.py b/.bin/Scripts/system_checklist.py
index 9169d1ff..38a82184 100644
--- a/.bin/Scripts/system_checklist.py
+++ b/.bin/Scripts/system_checklist.py
@@ -25,10 +25,11 @@ if __name__ == '__main__':
     ticket_number = get_ticket_number()
     other_results = {
       'Error': {
-        'BIOSKeyNotFoundError':   'BIOS key not found',
-        'CalledProcessError':     'Unknown Error',
-        'FileNotFoundError':    'File not found',
-        'GenericError':       'Unknown Error',
+        'BIOSKeyNotFoundError':     'BIOS key not found',
+        'CalledProcessError':       'Unknown Error',
+        'FileNotFoundError':        'File not found',
+        'GenericError':             'Unknown Error',
+        'Not4KAlignedError':        'False',
         'SecureBootDisabledError':  'Disabled',
       },
       'Warning': {
@@ -120,6 +121,9 @@ if __name__ == '__main__':
     try_and_print(message='Installed Office:',
       function=get_installed_office, ns='Unknown',
       other_results=other_results, print_return=True)
+    try_and_print(message='Partitions 4K aligned:',
+      function=check_4k_alignment, cs='True',
+      other_results=other_results)
     if D7_MODE:
       try_and_print(message='Product Keys:',
         function=get_product_keys, ns='Unknown', print_return=True)
diff --git a/.bin/Scripts/system_checklist_hw.py b/.bin/Scripts/system_checklist_hw.py
index 2f456003..b399064b 100644
--- a/.bin/Scripts/system_checklist_hw.py
+++ b/.bin/Scripts/system_checklist_hw.py
@@ -23,10 +23,11 @@ if __name__ == '__main__':
     ticket_number = get_ticket_number()
     other_results = {
       'Error': {
-        'BIOSKeyNotFoundError':   'BIOS key not found',
-        'CalledProcessError':     'Unknown Error',
-        'FileNotFoundError':    'File not found',
-        'GenericError':       'Unknown Error',
+        'BIOSKeyNotFoundError':     'BIOS key not found',
+        'CalledProcessError':       'Unknown Error',
+        'FileNotFoundError':        'File not found',
+        'GenericError':             'Unknown Error',
+        'Not4KAlignedError':        'False',
         'SecureBootDisabledError':  'Disabled',
       },
       'Warning': {
@@ -90,6 +91,9 @@ if __name__ == '__main__':
     try_and_print(message='Installed Office:',
       function=get_installed_office, ns='Unknown',
       other_results=other_results, print_return=True)
+    try_and_print(message='Partitions 4K aligned:',
+      function=check_4k_alignment, cs='True',
+      other_results=other_results)
 
     # Play audio, show devices, open Windows updates, and open Activation
     try_and_print(message='Opening Device Manager...',