From 2fe059925569a2b481e8d43b49c7b175b4682c99 Mon Sep 17 00:00:00 2001
From: rdb <git@rdb.name>
Date: Sat, 17 Aug 2019 22:42:52 +0200
Subject: [PATCH] deploy-ng: support adding icons to Windows binaries (part of
 #468)

---
 direct/src/dist/commands.py       |  47 +++++++
 direct/src/p3d/DeploymentTools.py | 214 ++++++++++++++++++++++++------
 2 files changed, 224 insertions(+), 37 deletions(-)

diff --git a/direct/src/dist/commands.py b/direct/src/dist/commands.py
index e5d9ac4f72..204fed4852 100644
--- a/direct/src/dist/commands.py
+++ b/direct/src/dist/commands.py
@@ -14,12 +14,14 @@ import struct
 import imp
 import string
 import time
+import tempfile
 
 import setuptools
 import distutils.log
 
 from . import FreezeTool
 from . import pefile
+from direct.p3d.DeploymentTools import Icon
 import panda3d.core as p3d
 
 
@@ -224,6 +226,7 @@ class build_apps(setuptools.Command):
         self.exclude_patterns = []
         self.include_modules = {}
         self.exclude_modules = {}
+        self.icons = {}
         self.platforms = [
             'manylinux1_x86_64',
             'macosx_10_6_x86_64',
@@ -298,6 +301,7 @@ class build_apps(setuptools.Command):
             key: _parse_list(value)
             for key, value in _parse_dict(self.exclude_modules).items()
         }
+        self.icons = _parse_dict(self.icons)
         self.platforms = _parse_list(self.platforms)
         self.plugins = _parse_list(self.plugins)
         self.extra_prc_files = _parse_list(self.extra_prc_files)
@@ -357,6 +361,18 @@ class build_apps(setuptools.Command):
         tmp.update(self.package_data_dirs)
         self.package_data_dirs = tmp
 
+        self.icon_objects = {}
+        for app, iconpaths in self.icons.items():
+            if not isinstance(iconpaths, list) and not isinstance(iconpaths, tuple):
+                iconpaths = (iconpaths,)
+
+            iconobj = Icon()
+            for iconpath in iconpaths:
+                iconobj.addImage(iconpath)
+
+            iconobj.generateMissingImages()
+            self.icon_objects[app] = iconobj
+
     def run(self):
         self.announce('Building platforms: {0}'.format(','.join(self.platforms)), distutils.log.INFO)
 
@@ -433,6 +449,22 @@ class build_apps(setuptools.Command):
 
         return wheelpaths
 
+    def update_pe_resources(self, appname, runtime):
+        """Update resources (e.g., icons) in windows PE file"""
+
+        icon = self.icon_objects.get(
+            appname,
+            self.icon_objects.get('*', None),
+        )
+
+        if icon is not None:
+            pef = pefile.PEFile()
+            pef.open(runtime, 'r+')
+            pef.add_icon(icon)
+            pef.add_resource_section()
+            pef.write_changes()
+            pef.close()
+
     def bundle_macos_app(self, builddir):
         """Bundle built runtime into a .app for macOS"""
 
@@ -618,6 +650,18 @@ class build_apps(setuptools.Command):
                 stub_path = os.path.join(os.path.dirname(dtool_path), '..', 'bin', stub_name)
                 stub_file = open(stub_path, 'rb')
 
+            # Do we need an icon?  On Windows, we need to add this to the stub
+            # before we add the blob.
+            if 'win' in platform:
+                temp_file = tempfile.NamedTemporaryFile(suffix='-icon.exe', delete=False)
+                temp_file.write(stub_file.read())
+                stub_file.close()
+                temp_file.close()
+                self.update_pe_resources(appname, temp_file.name)
+                stub_file = open(temp_file.name, 'rb')
+            else:
+                temp_file = None
+
             freezer.generateRuntimeFromStub(target_path, stub_file, use_console, {
                 'prc_data': prcexport if self.embed_prc_data else None,
                 'default_prc_dir': self.default_prc_dir,
@@ -633,6 +677,9 @@ class build_apps(setuptools.Command):
             }, self.log_append)
             stub_file.close()
 
+            if temp_file:
+                os.unlink(temp_file.name)
+
             # Copy the dependencies.
             search_path = [builddir]
             if use_wheels:
diff --git a/direct/src/p3d/DeploymentTools.py b/direct/src/p3d/DeploymentTools.py
index 9c390fbf07..70bf1da1b7 100644
--- a/direct/src/p3d/DeploymentTools.py
+++ b/direct/src/p3d/DeploymentTools.py
@@ -332,6 +332,132 @@ class Icon:
 
         return True
 
+    def generateMissingImages(self):
+        """ Generates image sizes that should be present but aren't by scaling
+        from the next higher size. """
+
+        for required_size in (48, 32, 24, 16):
+            if required_size in self.images:
+                continue
+
+            sizes = sorted(self.images.keys())
+            for from_size in sizes:
+                if from_size > required_size:
+                    break
+
+            if from_size > required_size:
+                Icon.notify.warning("Generating %dx%d icon by scaling down %dx%d image" % (required_size, required_size, from_size, from_size))
+
+                image = PNMImage(required_size, required_size)
+                if self.images[from_size].hasAlpha():
+                    image.addAlpha()
+                image.quickFilterFrom(self.images[from_size])
+                self.images[required_size] = image
+            else:
+                Icon.notify.warning("Cannot generate %dx%d icon; no higher resolution image available" % (required_size, required_size))
+
+    def _write_bitmap(self, fp, image, size, bpp):
+        """ Writes the bitmap header and data of an .ico file. """
+
+        fp.write(struct.pack('<IiiHHIIiiII', 40, size, size * 2, 1, bpp, 0, 0, 0, 0, 0, 0))
+
+        # XOR mask
+        if bpp == 24:
+            # Align rows to 4-byte boundary
+            rowalign = '\0' * (-(size * 3) & 3)
+            for y in xrange(size):
+                for x in xrange(size):
+                    r, g, b = image.getXel(x, size - y - 1)
+                    fp.write(struct.pack('<BBB', int(b * 255), int(g * 255), int(r * 255)))
+                fp.write(rowalign)
+
+        elif bpp == 32:
+            for y in xrange(size):
+                for x in xrange(size):
+                    r, g, b, a = image.getXelA(x, size - y - 1)
+                    fp.write(struct.pack('<BBBB', int(b * 255), int(g * 255), int(r * 255), int(a * 255)))
+
+        elif bpp == 8:
+            # We'll have to generate a palette of 256 colors.
+            hist = PNMImage.Histogram()
+            if image.hasAlpha():
+                # Make a copy without alpha channel.
+                image2 = PNMImage(image)
+                image2.premultiplyAlpha()
+                image2.removeAlpha()
+            else:
+                image2 = image
+            image2.make_histogram(hist)
+            colors = list(hist.get_pixels())
+            if len(colors) > 256:
+                # Palette too large; remove infrequent colors.
+                colors.sort(key=hist.get_count, reverse=True)
+
+                # Find the closest color on the palette matching each color
+                # that didn't fit.  This is certainly not the best palette
+                # generation code, but it'll do for now.
+                closest_indices = []
+                for color in colors[256:]:
+                    closest_index = 0
+                    closest_diff = 1025
+                    for i, closest_color in enumerate(colors[:256]):
+                        diff = abs(color.get_red() - closest_color.get_red()) \
+                             + abs(color.get_green() - closest_color.get_green()) \
+                             + abs(color.get_blue() - closest_color.get_blue())
+                        if diff < closest_diff:
+                            closest_index = i
+                            closest_diff = diff
+                    assert closest_diff < 100
+                    closest_indices.append(closest_index)
+
+            # Write the palette.
+            i = 0
+            while i < 256 and i < len(colors):
+                r, g, b, a = colors[i]
+                fp.write(struct.pack('<BBBB', b, g, r, 0))
+                i += 1
+            if i < 256:
+                # Fill the rest with zeroes.
+                fp.write(b'\x00' * (4 * (256 - i)))
+
+            # Write indices.  Align rows to 4-byte boundary.
+            rowalign = b'\0' * (-size & 3)
+            for y in xrange(size):
+                for x in xrange(size):
+                    pixel = image2.get_pixel(x, size - y - 1)
+                    index = colors.index(pixel)
+                    if index >= 256:
+                        # Find closest pixel instead.
+                        index = closest_indices[index - 256]
+                    fp.write(struct.pack('<B', index))
+                fp.write(rowalign)
+        else:
+            raise ValueError("Invalid bpp %d" % (bpp))
+
+        # Create an AND mask, aligned to 4-byte boundary
+        if image.hasAlpha() and bpp <= 8:
+            rowalign = b'\0' * (-((size + 7) >> 3) & 3)
+            for y in xrange(size):
+                mask = 0
+                num_bits = 7
+                for x in xrange(size):
+                    a = image.get_alpha_val(x, size - y - 1)
+                    if a <= 1:
+                        mask |= (1 << num_bits)
+                    num_bits -= 1
+                    if num_bits < 0:
+                        fp.write(struct.pack('<B', mask))
+                        mask = 0
+                        num_bits = 7
+                if num_bits < 7:
+                    fp.write(struct.pack('<B', mask))
+                fp.write(rowalign)
+        else:
+            andsize = (size + 7) >> 3
+            if andsize % 4 != 0:
+                andsize += 4 - (andsize % 4)
+            fp.write(b'\x00' * (andsize * size))
+
     def makeICO(self, fn):
         """ Writes the images to a Windows ICO file.  Returns True on success. """
 
@@ -339,57 +465,71 @@ class Icon:
             fn = Filename.fromOsSpecific(fn)
         fn.setBinary()
 
+        # ICO files only support resolutions up to 256x256.
         count = 0
         for size in self.images.keys():
+            if size < 256:
+                count += 1
             if size <= 256:
                 count += 1
+        dataoffs = 6 + count * 16
 
         ico = open(fn, 'wb')
         ico.write(struct.pack('<HHH', 0, 1, count))
 
-        # Write the directory
+        # Write 8-bpp image headers for sizes under 256x256.
         for size, image in self.images.items():
-            if size == 256:
+            if size >= 256:
+                continue
+            ico.write(struct.pack('<BB', size, size))
+
+            # Calculate row sizes
+            xorsize = size
+            if xorsize % 4 != 0:
+                xorsize += 4 - (xorsize % 4)
+            andsize = (size + 7) >> 3
+            if andsize % 4 != 0:
+                andsize += 4 - (andsize % 4)
+            datasize = 40 + 256 * 4 + (xorsize + andsize) * size
+
+            ico.write(struct.pack('<BBHHII', 0, 0, 1, 8, datasize, dataoffs))
+            dataoffs += datasize
+
+        # Write 24/32-bpp image headers.
+        for size, image in self.images.items():
+            if size > 256:
+                continue
+            elif size == 256:
                 ico.write('\0\0')
             else:
                 ico.write(struct.pack('<BB', size, size))
-            bpp = 32 if image.hasAlpha() else 24
-            ico.write(struct.pack('<BBHHII', 0, 0, 1, bpp, 0, 0))
 
-        # Now write the actual icons
-        ptr = 14
-        for size, image in self.images.items():
-            loc = ico.tell()
-            bpp = 32 if image.hasAlpha() else 24
-            ico.write(struct.pack('<IiiHHIIiiII', 40, size, size * 2, 1, bpp, 0, 0, 0, 0, 0, 0))
-
-            # XOR mask
-            if bpp == 24:
-                # Align rows to 4-byte boundary
-                rowalign = '\0' * (-(size * 3) & 3)
-                for y in xrange(size):
-                    for x in xrange(size):
-                        r, g, b = image.getXel(x, size - y - 1)
-                        ico.write(struct.pack('<BBB', int(b * 255), int(g * 255), int(r * 255)))
-                    ico.write(rowalign)
+            # Calculate the size so we can write the offset within the file.
+            if image.hasAlpha():
+                bpp = 32
+                xorsize = size * 4
             else:
-                for y in xrange(size):
-                    for x in xrange(size):
-                        r, g, b, a = image.getXelA(x, size - y - 1)
-                        ico.write(struct.pack('<BBBB', int(b * 255), int(g * 255), int(r * 255), int(a * 255)))
+                bpp = 24
+                xorsize = size * 3 + (-(size * 3) & 3)
+            andsize = (size + 7) >> 3
+            if andsize % 4 != 0:
+                andsize += 4 - (andsize % 4)
+            datasize = 40 + (xorsize + andsize) * size
 
-            # Empty AND mask, aligned to 4-byte boundary
-            #TODO: perhaps we should convert alpha into an AND mask
-            # to support older versions of Windows that don't support alpha.
-            ico.write('\0' * (size * (size / 8 + (-((size / 8) * 3) & 3))))
+            ico.write(struct.pack('<BBHHII', 0, 0, 1, bpp, datasize, dataoffs))
+            dataoffs += datasize
 
-            # Go back to write the location
-            dataend = ico.tell()
-            ico.seek(ptr)
-            ico.write(struct.pack('<II', dataend - loc, loc))
-            ico.seek(dataend)
-            ptr += 16
+        # Now write the actual icon bitmap data.
+        for size, image in self.images.items():
+            if size < 256:
+                self._write_bitmap(ico, image, size, 8)
 
+        for size, image in self.images.items():
+            if size <= 256:
+                bpp = 32 if image.hasAlpha() else 24
+                self._write_bitmap(ico, image, size, bpp)
+
+        assert ico.tell() == dataoffs
         ico.close()
 
         return True
@@ -406,9 +546,9 @@ class Icon:
         icns = open(stream, 'wb')
         icns.write(b'icns\0\0\0\0')
 
-        icon_types = {16: 'is32', 32: 'il32', 48: 'ih32', 128: 'it32'}
-        mask_types = {16: 's8mk', 32: 'l8mk', 48: 'h8mk', 128: 't8mk'}
-        png_types = {256: 'ic08', 512: 'ic09'}
+        icon_types = {16: b'is32', 32: b'il32', 48: b'ih32', 128: b'it32'}
+        mask_types = {16: b's8mk', 32: b'l8mk', 48: b'h8mk', 128: b't8mk'}
+        png_types = {256: b'ic08', 512: b'ic09'}
 
         pngtype = PNMFileTypeRegistry.getGlobalPtr().getTypeFromExtension("png")