From f1645f80e6d5544ed23dc0a7d4cfab4de338c7a4 Mon Sep 17 00:00:00 2001
From: 2Shirt <2xShirt@gmail.com>
Date: Wed, 29 Sep 2021 20:27:35 -0600
Subject: [PATCH] Add AV check to Auto Setup

---
 scripts/auto_setup.py   |  2 +-
 scripts/wk/os/win.py    | 41 +++++++++++++++++++++++++++++++++++++++++
 scripts/wk/setup/win.py |  8 +++++++-
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/scripts/auto_setup.py b/scripts/auto_setup.py
index 3728c53d..2d75585c 100644
--- a/scripts/auto_setup.py
+++ b/scripts/auto_setup.py
@@ -142,7 +142,7 @@ BASE_MENUS = {
       MenuEntry('Secure Boot',              'auto_show_secure_boot_status'),
       MenuEntry('Installed RAM',            'auto_show_installed_ram'),
       MenuEntry('Storage Status',           'auto_show_storage_status'),
-      MenuEntry('Virus Protection',         no_op),
+      MenuEntry('Virus Protection',         'auto_show_installed_antivirus'),
       MenuEntry('Partitions 4K Aligned',    no_op),
       ),
     'Run Programs': (
diff --git a/scripts/wk/os/win.py b/scripts/wk/os/win.py
index e3a7bc7c..50ea654e 100644
--- a/scripts/wk/os/win.py
+++ b/scripts/wk/os/win.py
@@ -152,6 +152,47 @@ def set_timezone(zone):
 
 
 # Info Functions
+def get_installed_antivirus():
+  """Get list of installed antivirus programs, returns list."""
+  cmd = [
+    'WMIC', r'/namespace:\\root\SecurityCenter2',
+    'path', 'AntivirusProduct',
+    'get', 'displayName', '/value',
+    ]
+  products = []
+  report = []
+
+  # Get list of products
+  proc = run_program(cmd)
+  for line in proc.stdout.splitlines():
+    line = line.strip()
+    if '=' in line:
+      products.append(line.split('=')[1])
+
+  # Check product(s) status
+  for product in sorted(products):
+    cmd = [
+      'WMIC', r'/namespace:\\root\SecurityCenter2',
+      'path', 'AntivirusProduct',
+      'where', f'displayName="{product}"',
+      'get', 'productState', '/value',
+      ]
+    proc = run_program(cmd)
+    state = proc.stdout.split('=')[1]
+    state = hex(int(state))
+    if str(state)[3:5] not in ['10', '11']:
+      report.append(color_string(f'[Disabled] {product}', 'YELLOW'))
+    else:
+      report.append(product)
+
+  # Final check
+  if not report:
+    report.append(color_string('No products detected', 'RED'))
+
+  # Done
+  return report
+
+
 def get_installed_ram(as_list=False, raise_exceptions=False):
   """Get installed RAM."""
   mem = psutil.virtual_memory()
diff --git a/scripts/wk/setup/win.py b/scripts/wk/setup/win.py
index a1c75c6d..a10b2ec0 100644
--- a/scripts/wk/setup/win.py
+++ b/scripts/wk/setup/win.py
@@ -44,13 +44,13 @@ if platform.system() == 'Windows':
   from wk.os.win        import (
     OS_VERSION,
     activate_with_bios,
+    get_installed_antivirus,
     get_installed_ram,
     get_os_activation,
     get_os_name,
     get_raw_disks,
     get_volume_usage,
     is_secure_boot_enabled,
-    reg_read_value,
     reg_set_value,
     reg_write_settings,
     )
@@ -71,6 +71,7 @@ else:
     """No-op function."""
   # wk.os.win
   activate_with_bios = no_op
+  get_installed_antivirus = no_op
   get_installed_ram = no_op
   get_os_activation = no_op
   get_os_name = no_op
@@ -569,6 +570,11 @@ def auto_restore_default_uac():
   TRY_PRINT.run('User Account Control...', restore_default_uac)
 
 
+def auto_show_installed_antivirus():
+  """Display installed antivirus."""
+  TRY_PRINT.run('Virus Protection...', get_installed_antivirus)
+
+
 def auto_show_installed_ram():
   """Display installed RAM."""
   TRY_PRINT.run('Installed RAM...', get_installed_ram,