From 15922ac706f2839177dcd320cc39104639c78148 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Pawe=C5=82=20Jastrz=C4=99bski?= <pawelj@iosphe.re>
Date: Thu, 24 Oct 2019 14:48:54 +0200
Subject: [PATCH] Added symlink_protection command

---
 CB/IconChanger.py | 67 +++++++++++++++++++++++++++++++++++++++++++++++
 CB/__init__.py    |  2 +-
 CurseBreaker.py   | 13 +++++++++
 README.md         |  1 +
 4 files changed, 82 insertions(+), 1 deletion(-)
 create mode 100644 CB/IconChanger.py

diff --git a/CB/IconChanger.py b/CB/IconChanger.py
new file mode 100644
index 0000000..ac8aef1
--- /dev/null
+++ b/CB/IconChanger.py
@@ -0,0 +1,67 @@
+import os
+import ctypes
+from ctypes import POINTER, Structure, c_wchar, c_int, sizeof, byref
+from ctypes.wintypes import BYTE, WORD, DWORD, LPWSTR
+
+HICON = c_int
+LPTSTR = LPWSTR
+TCHAR = c_wchar
+MAX_PATH = 260
+FCSM_ICONFILE = 0x00000010
+FCS_FORCEWRITE = 0x00000002
+SHGFI_ICONLOCATION = 0x000001000
+
+
+class GUID(Structure):
+    _fields_ = [
+        ('Data1', DWORD),
+        ('Data2', WORD),
+        ('Data3', WORD),
+        ('Data4', BYTE * 8)]
+
+
+class SHFOLDERCUSTOMSETTINGS(Structure):
+    _fields_ = [
+        ('dwSize', DWORD),
+        ('dwMask', DWORD),
+        ('pvid', POINTER(GUID)),
+        ('pszWebViewTemplate', LPTSTR),
+        ('cchWebViewTemplate', DWORD),
+        ('pszWebViewTemplateVersion', LPTSTR),
+        ('pszInfoTip', LPTSTR),
+        ('cchInfoTip', DWORD),
+        ('pclsid', POINTER(GUID)),
+        ('dwFlags', DWORD),
+        ('pszIconFile', LPTSTR),
+        ('cchIconFile', DWORD),
+        ('iIconIndex', c_int),
+        ('pszLogo', LPTSTR),
+        ('cchLogo', DWORD)]
+
+
+class SHFILEINFO(Structure):
+    _fields_ = [
+        ('hIcon', HICON),
+        ('iIcon', c_int),
+        ('dwAttributes', DWORD),
+        ('szDisplayName', TCHAR * MAX_PATH),
+        ('szTypeName', TCHAR * 80)]
+
+
+def set_icon(folderpath, iconpath, iconindex):
+    folderpath = os.path.abspath(folderpath)
+    iconpath = os.path.abspath(iconpath)
+
+    fcs = SHFOLDERCUSTOMSETTINGS()
+    fcs.dwSize = sizeof(fcs)
+    fcs.dwMask = FCSM_ICONFILE
+    fcs.pszIconFile = iconpath
+    fcs.cchIconFile = 0
+    fcs.iIconIndex = iconindex
+
+    ctypes.windll.shell32.SHGetSetFolderCustomSettings(byref(fcs), folderpath, FCS_FORCEWRITE)
+    sfi = SHFILEINFO()
+    ctypes.windll.shell32.SHGetFileInfoW(folderpath, 0, byref(sfi), sizeof(sfi), SHGFI_ICONLOCATION)
+    index = ctypes.windll.shell32.Shell_GetCachedImageIndexW(sfi.szDisplayName, sfi.iIcon, 0)
+    ctypes.windll.shell32.SHUpdateImageW(sfi.szDisplayName, sfi.iIcon, 0, index)
+
diff --git a/CB/__init__.py b/CB/__init__.py
index df5c89f..7733d8f 100644
--- a/CB/__init__.py
+++ b/CB/__init__.py
@@ -1,7 +1,7 @@
 import string
 import random
 
-__version__ = '3.0.2'
+__version__ = '3.1.0'
 __license__ = 'GPLv3'
 __copyright__ = '2019, Paweł Jastrzębski <pawelj@iosphe.re>'
 __docformat__ = 'restructuredtext en'
diff --git a/CurseBreaker.py b/CurseBreaker.py
index 724d1ee..69ce4d9 100644
--- a/CurseBreaker.py
+++ b/CurseBreaker.py
@@ -23,6 +23,7 @@
 
 if platform.system() == 'Windows':
     from ctypes import windll, wintypes, byref
+    from CB.IconChanger import set_icon
 
 
 class TUI:
@@ -389,6 +390,18 @@ def c_uri_integration(self, _):
         else:
             printft('This feature is available only on Windows.')
 
+    def c_symlink_protection(self, _):
+        if self.os == 'Windows':
+            printft(HTML('<ansigreen>Directories tweaked:</ansigreen>'))
+            for root, dirs, _ in os.walk(self.core.path / '..' / '..'):
+                for d in dirs:
+                    path = Path(root) / d
+                    if os.path.islink(path):
+                        set_icon(path, Path("C:/Windows/System32/SHELL32.dll"), 4)
+                        print(path.resolve())
+        else:
+            printft('This feature is available only on Windows.')
+
     def c_toggle_dev(self, args):
         if args:
             status = self.core.dev_toggle(args)
diff --git a/README.md b/README.md
index c354bd7..addb78d 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,7 @@ By default **CurseBreaker** will create backups of entire `WTF` directory.
 - Most of the commands support the comma-separated list of addons.
 - `install` command have optional `-i` flag that can be used to disable client version check.
 - Environment variable `CURSEBREAKER_PATH` can be used to set the custom location of WoW client.
+- Undocumented `symlink_protection` command will protect your symlinks against endless hunger of Battle.NET client.
 
 ## SUPPORTED URL
 - CurseForge: `https://www.curseforge.com/wow/addons/[addon_name]`, `cf:[addon_name]`