From 7816602685b1c36c4ee3aae9fedef47867f87571 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Tue, 4 Jun 2019 18:40:42 -0600
Subject: [PATCH] 4K Alignment checks

---
 .bin/Scripts/functions/common.py   |  7 +++++--
 .bin/Scripts/functions/hw_diags.py | 26 +++++++++++++++++++++++++-
 .bin/Scripts/functions/sw_diags.py | 29 +++++++++++++++++++++++++++++
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/.bin/Scripts/functions/common.py b/.bin/Scripts/functions/common.py
index 28834660..9db75814 100644
--- a/.bin/Scripts/functions/common.py
+++ b/.bin/Scripts/functions/common.py
@@ -64,10 +64,13 @@ class GenericRepair(Exception):
 class MultipleInstallationsError(Exception):
   pass
 
-class NotInstalledError(Exception):
+class NoProfilesError(Exception):
   pass
 
-class NoProfilesError(Exception):
+class Not4KAlignedError(Exception):
+  pass
+
+class NotInstalledError(Exception):
   pass
 
 class OSInstalledLegacyError(Exception):
diff --git a/.bin/Scripts/functions/hw_diags.py b/.bin/Scripts/functions/hw_diags.py
index c085a5ed..36be4a30 100644
--- a/.bin/Scripts/functions/hw_diags.py
+++ b/.bin/Scripts/functions/hw_diags.py
@@ -307,6 +307,11 @@ class DiskObj():
         attr_type=self.attr_type, **COLORS))
       report.extend(sorted(self.nvme_smart_notes.keys()))
 
+    # 4K alignment check
+    if not self.is_4k_aligned():
+      report.append('{YELLOW}Warning{CLEAR}'.format(**COLORS))
+      report.append('  One or more partitions are not 4K aligned')
+
     # Tests
     for test in self.tests.values():
       report.extend(test.report)
@@ -410,6 +415,26 @@ class DiskObj():
             'self_test', {}).get(
             k, {})
 
+  def is_4k_aligned(self):
+    """Check if partitions are 4K aligned, returns bool."""
+    cmd = [
+      'sudo',
+      'sfdisk',
+      '--json',
+      self.path,
+      ]
+    aligned = True
+
+    # Get partition details
+    json_data = get_json_from_command(cmd)
+
+    # Check partitions
+    for part in json_data.get('partitiontable', {}).get('partitions', []):
+      aligned = aligned and part.get('start', -1) % 4096 == 0
+
+    # Done
+    return aligned
+
   def safety_check(self, silent=False):
     """Run safety checks and disable tests if necessary."""
     test_running = False
@@ -454,7 +479,6 @@ class DiskObj():
         disk_ok = OVERRIDES_FORCED or ask('Run tests on this device anyway?')
         print_standard(' ')
 
-
     # Disable tests if necessary (statuses won't be overwritten)
     if test_running:
       if not silent:
diff --git a/.bin/Scripts/functions/sw_diags.py b/.bin/Scripts/functions/sw_diags.py
index 1b965766..b122528a 100644
--- a/.bin/Scripts/functions/sw_diags.py
+++ b/.bin/Scripts/functions/sw_diags.py
@@ -6,6 +6,35 @@ from functions.common import *
 from settings.sw_diags import *
 
 
+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 = aligned and 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: