first upload

setup.py

@@ -1,78 +1,193 @@
-
 # -*- coding: utf-8 -*-
+"""WAPT agent-side configuration of cleanmgr.exe profiles.
+
+The WAPT update_package.py wizard generates profiles.json in the package sources.
+Installing this package configures VolumeCaches only: it DOES NOT run cleanup.
+"""
+
 from setuphelpers import *
+import json
+import re
 import winreg
 
 uninstallkey = []
 
-PROFILE = 100
-REG_PATH = (
-    r'SOFTWARE\Microsoft\Windows\CurrentVersion'
-    r'\Explorer\VolumeCaches'
-)
-
-CATEGORIES = {
-    'Temporary Files',
-    'Thumbnail Cache',
-    'Delivery Optimization Files',
-}
-
-FLAG = 'StateFlags{:04d}'.format(PROFILE)
+VOLUME_CACHES = r'SOFTWARE\Microsoft\Windows\CurrentVersion\Explorer\VolumeCaches'
+TRACK_KEY = r'SOFTWARE\Comitari\CleanmgrProfiles'
+TRACK_VALUE = 'ProfilesJson'
+PROFILE_LIMIT = 50
 
 
-def get_categories():
-    with reg_openkey_noredir(
-        winreg.HKEY_LOCAL_MACHINE, REG_PATH
-    ) as root:
-        categories = []
-        index = 0
+def _validate_config(config):
+    if not isinstance(config, dict) or config.get('schema_version') != 1:
+        raise ValueError('Invalid profiles.json schema_version')
+    profiles = config.get('profiles')
+    if not isinstance(profiles, list) or not 1 <= len(profiles) <= PROFILE_LIMIT:
+        raise ValueError('Expected 1..%d profiles' % PROFILE_LIMIT)
+
+    seen = set()
+    for profile in profiles:
+        if not isinstance(profile, dict):
+            raise ValueError('Invalid profile')
+        profile_id = profile.get('id')
+        if type(profile_id) is not int or not 0 <= profile_id <= 9999:
+            raise ValueError('Profile id must be in 0..9999')
+        if profile_id in seen:
+            raise ValueError('Duplicate profile id: %d' % profile_id)
+        seen.add(profile_id)
+
+        name = profile.get('name')
+        if not isinstance(name, str) or not 1 <= len(name.strip()) <= 80:
+            raise ValueError('Profile name missing or too long')
+        categories = profile.get('categories')
+        if not isinstance(categories, list) or not categories:
+            raise ValueError('Profile %d has no categories' % profile_id)
+        for category in categories:
+            if not isinstance(category, str) or not re.fullmatch(
+                    r'[A-Za-z0-9][A-Za-z0-9 ()_.-]{0,127}', category):
+                raise ValueError('Unsafe or invalid VolumeCaches key: %r' % category)
+        if len(set(categories)) != len(categories):
+            raise ValueError('Duplicate categories in profile %d' % profile_id)
+    return config
+
+
+def _read_deployed_config():
+    try:
+        with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, TRACK_KEY,
+                                sam=winreg.KEY_READ) as key:
+            raw = winreg.QueryValueEx(key, TRACK_VALUE)[0]
+    except FileNotFoundError:
+        return None
+    return _validate_config(json.loads(raw))
+
+
+def _store_deployed_config(config):
+    with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, TRACK_KEY,
+                            sam=winreg.KEY_WRITE, create_if_missing=True) as key:
+        reg_setvalue(key, TRACK_VALUE,
+                     json.dumps(config, ensure_ascii=False, sort_keys=True),
+                     type=winreg.REG_SZ)
+
+
+def _list_volume_categories():
+    with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, VOLUME_CACHES,
+                            sam=winreg.KEY_READ) as root:
+        result = []
+        position = 0
         while True:
             try:
-                categories.append(winreg.EnumKey(root, index))
-                index += 1
+                result.append(winreg.EnumKey(root, position))
             except OSError:
                 break
-    return categories
+            position += 1
+    return sorted(result)
+
+
+def _flag(profile_id):
+    return 'StateFlags%04d' % profile_id
+
+
+def _delete_profile_flags(profile_id, categories):
+    for category in categories:
+        path = VOLUME_CACHES + '\\' + category
+        with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, path,
+                                sam=winreg.KEY_SET_VALUE) as key:
+            try:
+                winreg.DeleteValue(key, _flag(profile_id))
+            except FileNotFoundError:
+                pass
+
+
+def _apply_profile(profile, available):
+    wanted = set(profile['categories'])
+    flag = _flag(profile['id'])
+    for category in available:
+        path = VOLUME_CACHES + '\\' + category
+        with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, path,
+                                sam=winreg.KEY_SET_VALUE) as key:
+            reg_setvalue(key, flag, 2 if category in wanted else 0,
+                         type=winreg.REG_DWORD)
+    unavailable = wanted - set(available)
+    if unavailable:
+        print('Profile %d (%s), missing on this Windows: %s' % (
+            profile['id'], profile['name'], ', '.join(sorted(unavailable))))
+    print('Configured profile %d (%s): %d available selected categories' % (
+        profile['id'], profile['name'], len(wanted & set(available))))
 
 
 def install():
-    available = get_categories()
+    path = makepath(basedir, 'profiles.json')
+    with open(path, 'r', encoding='utf-8') as stream:
+        config = _validate_config(json.load(stream))
 
-    for category in available:
-        path = REG_PATH + '\\' + category
+    available = _list_volume_categories()
+    if not available:
+        raise RuntimeError('No VolumeCaches categories available on this Windows')
 
-        with reg_openkey_noredir(
-            winreg.HKEY_LOCAL_MACHINE,
-            path,
-            sam=winreg.KEY_SET_VALUE
-        ) as key:
-            value = 2 if category in CATEGORIES else 0
-            reg_setvalue(
-                key, FLAG, value, type=winreg.REG_DWORD
-            )
+    previous = _read_deployed_config()
+    previous_ids = {p['id'] for p in previous['profiles']} if previous else set()
+    new_ids = {p['id'] for p in config['profiles']}
+    for removed_id in sorted(previous_ids - new_ids):
+        _delete_profile_flags(removed_id, available)
+        print('Removed former profile %d' % removed_id)
 
-        if category in CATEGORIES:
-            print('Enabled: {}'.format(category))
+    for profile in config['profiles']:
+        _apply_profile(profile, available)
 
-    missing = CATEGORIES - set(available)
-    if missing:
-        print('Unavailable categories: {}'.format(missing))
+    _store_deployed_config(config)
+    print('Cleanmgr profiles configured. No cleaning was executed.')
 
-    # Execution silencieuse du profil
-    run('cleanmgr.exe /sagerun:{}'.format(PROFILE),
-        timeout=3600)
+
+def audit():
+    try:
+        config = _read_deployed_config()
+        if not config:
+            print('Cleanmgr profile configuration missing')
+            return 'ERROR'
+        available = _list_volume_categories()
+        if not available:
+            return 'ERROR'
+        missing = False
+        errors = []
+        for profile in config['profiles']:
+            wanted = set(profile['categories'])
+            missing_categories = wanted - set(available)
+            if missing_categories:
+                print('Profile %d unavailable categories: %s' % (
+                    profile['id'], ', '.join(sorted(missing_categories))))
+                missing = True
+            for category in available:
+                path = VOLUME_CACHES + '\\' + category
+                with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, path,
+                                        sam=winreg.KEY_READ) as key:
+                    try:
+                        current, kind = winreg.QueryValueEx(key, _flag(profile['id']))
+                    except FileNotFoundError:
+                        current, kind = None, None
+                expected = 2 if category in wanted else 0
+                if current != expected or kind != winreg.REG_DWORD:
+                    errors.append('%d / %s: expected %d, got %r' % (
+                        profile['id'], category, expected, current))
+        if errors:
+            for item in errors[:25]:
+                print(item)
+            return 'ERROR'
+        return 'WARNING' if missing else 'OK'
+    except Exception as exc:
+        print('Cleanmgr audit failed: %s' % exc)
+        return 'ERROR'
 
 
 def uninstall():
-    for category in get_categories():
-        path = REG_PATH + '\\' + category
-
-        with reg_openkey_noredir(
-            winreg.HKEY_LOCAL_MACHINE,
-            path,
-            sam=winreg.KEY_SET_VALUE
-        ) as key:
+    previous = _read_deployed_config()
+    if previous:
+        available = _list_volume_categories()
+        for profile in previous['profiles']:
+            _delete_profile_flags(profile['id'], available)
+        with reg_openkey_noredir(winreg.HKEY_LOCAL_MACHINE, TRACK_KEY,
+                                sam=winreg.KEY_SET_VALUE) as key:
             try:
-                winreg.DeleteValue(key, FLAG)
+                winreg.DeleteValue(key, TRACK_VALUE)
             except FileNotFoundError:
                 pass
+    print('Managed Cleanmgr profile flags removed. No cleaning was executed.')