From 526f6e26ebcdd43939b5a527cb35587cadc10f08 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Wed, 29 Sep 2021 01:24:07 -0600
Subject: [PATCH] Add installed RAM sections

---
 scripts/auto_setup.py   |  2 +-
 scripts/wk/os/win.py    | 20 +++++++++++++++++++-
 scripts/wk/setup/win.py | 11 ++++++++++-
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/scripts/auto_setup.py b/scripts/auto_setup.py
index 5d614f33..7d2c1f42 100644
--- a/scripts/auto_setup.py
+++ b/scripts/auto_setup.py
@@ -140,7 +140,7 @@ BASE_MENUS = {
       MenuEntry('Operating System',         'auto_show_os_name'),
       MenuEntry('Windows Activation',       'auto_show_activation'),
       MenuEntry('Secure Boot',              'auto_show_secure_boot_status'),
-      MenuEntry('Installed RAM',            no_op),
+      MenuEntry('Installed RAM',            'auto_show_installed_ram'),
       MenuEntry('Storage Volumes',          no_op),
       MenuEntry('Virus Protection',         no_op),
       MenuEntry('Partitions 4K Aligned',    no_op),
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index 41b52a26..b5efeaa7 100644
--- a/scripts/wk/os/win.py
+++ b/scripts/wk/os/win.py
@@ -24,7 +24,7 @@ from wk.cfg.windows_builds import (
   WINDOWS_BUILDS,
   )
 from wk.exe import run_program
-from wk.std import GenericError, GenericWarning, sleep
+from wk.std import GenericError, GenericWarning, bytes_to_string, sleep
 
 
 # STATIC VARIABLES
@@ -65,6 +65,8 @@ KNOWN_HIVE_NAMES = {
   winreg.HKEY_USERS: 'HKEY_USERS',
   }
 OS_VERSION = float(platform.win32_ver()[0])
+RAM_OK      = 5.5 * 1024**3 # ~6 GiB assuming a bit of shared memory
+RAM_WARNING = 3.5 * 1024**3 # ~4 GiB assuming a bit of shared memory
 REG_MSISERVER = r'HKLM\SYSTEM\CurrentControlSet\Control\SafeBoot\Network\MSIServer'
 SLMGR = pathlib.Path(f'{os.environ.get("SYSTEMROOT")}/System32/slmgr.vbs')
 
@@ -143,6 +145,22 @@ def set_timezone(zone):
 
 
 # Info Functions
+def get_installed_ram(as_list=False, raise_exceptions=False):
+  """Get installed RAM."""
+  mem = psutil.virtual_memory()
+  mem_str = bytes_to_string(mem.total, decimals=1)
+
+  # Raise exception if necessary
+  if raise_exceptions:
+    if RAM_OK > mem.total >= RAM_WARNING:
+      raise GenericWarning(mem_str)
+    if mem.total > RAM_WARNING:
+      raise GenericError(mem_str)
+
+  # Done
+  return [mem_str] if as_list else mem_str
+
+
 def get_os_activation(as_list=False, check=True):
   """Get OS activation status, returns str.
 
diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py
index c0d9f6b2..acda54b1 100644
--- a/scripts/wk/setup/win.py
+++ b/scripts/wk/setup/win.py
@@ -44,6 +44,7 @@ if platform.system() == 'Windows':
   from wk.os.win        import (
     OS_VERSION,
     activate_with_bios,
+    get_installed_ram,
     get_os_activation,
     get_os_name,
     is_secure_boot_enabled,
@@ -68,6 +69,7 @@ else:
     """No-op function."""
   # wk.os.win
   activate_with_bios = no_op
+  get_installed_ram = no_op
   get_os_activation = no_op
   get_os_name = no_op
   is_secure_boot_enabled = no_op
@@ -563,8 +565,15 @@ def auto_restore_default_uac():
   TRY_PRINT.run('User Account Control...', restore_default_uac)
 
 
+def auto_show_installed_ram():
+  """Display installed RAM."""
+  TRY_PRINT.run('Installed RAM...', get_installed_ram,
+  as_list=True, raise_exceptions=True,
+  )
+
+
 def auto_show_os_activation():
-  """Display OS Name."""
+  """Display OS activation status."""
   TRY_PRINT.run('Activation...', get_os_activation, as_list=True)