'''
.       .1111...          | Title: png2ico.py
    .10000000000011.   .. | Author: Oliver Morton (Sec-1 Ltd)
 .00              000...  | Email: oliverm-tools@sec-1.com
1                  01..   | Description:
                    ..    | Create a multisize ICO from a PNG file.
                   ..     | Requires:
GrimHacker        ..      | pillow
                 ..       |
grimhacker.com  ..        |
@grimhacker    ..         |
----------------------------------------------------------------------------
PNG2ICO - Create a multisize ICO from a PNG
    Copyright (C) 2015  Oliver Morton (Sec-1 Ltd)
    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License along
    with this program; if not, write to the Free Software Foundation, Inc.,
    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
'''
# Adapted from https://github.com/aclements/commuter/blob/master/web/png2ico
# and https://code.google.com/p/reply-manager/source/browse/build/png2ico.py

__version__ = "$Revision: 1.0 $"
# $Source$

import os
import struct
import ctypes
import logging
import argparse

from PIL import Image
from io import BytesIO


# ref: http://msdn.microsoft.com/en-us/library/ms997538
class ICONDIR(ctypes.LittleEndianStructure):
    _pack_ = 1
    _fields_ = [("idReserved", ctypes.c_ushort),
                ("idType", ctypes.c_ushort),
                ("idCount", ctypes.c_ushort)]


class ICONDIRENTRY(ctypes.LittleEndianStructure):
    _pack_ = 1
    _fields_ = [("bWidth", ctypes.c_byte),
                ("bHeight", ctypes.c_byte),
                ("bColorCount", ctypes.c_byte),
                ("bReserved", ctypes.c_byte),
                ("wPlanes", ctypes.c_ushort),
                ("wBitCount", ctypes.c_ushort),
                ("dwBytesInRes", ctypes.c_ulong),
                ("dwImageOffset", ctypes.c_ulong)]


class PNG2ICO(object):
    """PNG2ICO - Create a multisize ICO from a PNG
    First creates all required sizes of supplied image in memory then
    builds ICO file structure using struct and writes to specified
    output file."""
    def __init__(self, image_files, out_file, sizes=[256,48,32,16]):
        self.log = logging.getLogger(__name__)
        if image_files:
            self._image_files = image_files
        else:
            raise Exception("PNG image_files required")
        if out_file:
            self._out_file = out_file
        else:
            raise Exception("out_file required")
        if sizes and isinstance(sizes, list) and all(isinstance(item, int) for item in sizes):
            self._sizes = sorted(sizes, reverse=True)  # need largest first for ico format
        else:
            raise Exception("list of integer sizes required")

    @property
    def image_files(self):
        return self._image_files

    @property
    def out_file(self):
        return self._out_file

    @property
    def sizes(self):
        return self._sizes

    def _resize(self, infile, size):
        """Resize image to size x size.
        Returns file like object."""
        self.log.debug("Resizing {file_} to {size} x {size}".format(file_=infile, size=size))
        outfile = BytesIO()
        #outfile = "{0}_{1}.png".format(os.path.split(infile)[1], size)
        try:
            im = Image.open(infile).convert("RGBA")         # open and convert to RGBA
            im.thumbnail((size, size), Image.ANTIALIAS)     # Create thumbnail of designated size
            im.save(outfile, "PNG")                         # Save thumbnail to BytesIO instance
            outfile.seek(0)                                 # Rewind the BytesIO object to the start ready for reading.
            outfile.name = "{0}_{1}.png".format(os.path.split(infile)[1], size)  # Give the BytesIO object a name
        except Exception as e:
            raise Exception("Cannot resize '{0}' to {1}x{1} {2}".format(infile, size, e))
        else:
            #return open(outfile, "rb")
            return outfile

    def _create_ico(self, out_file, image_files):        
        """Create ico file from list of image file like objects.
        This builds the ico file using struct based on the specification."""
        self.log.debug("Loading image files...")
        images = []
        for image_file in image_files:
            data = image_file.read()
            # Check image_file is a PNG
            if not (data.startswith('\x89PNG\x0d\x0a\x1a\x0a') and
                data[12:].startswith('IHDR') and
                len(data) >= 29):
                    raise Exception('not a valid png file: {0}'.format(image_file.name))
            else:
                self.log.debug("{0} is a PNG image".format(image_file.name))
            w, h, bpp = struct.unpack('!IIB', data[16:16+9])  # Extract width, height, and bits per pixel
            if w > 256 or h > 256:  # Check maximum icon size is not exceeded
                raise Exception("{0} exceeds 256x256 size limit!".format(image_file.name))
            else:
                self.log.debug("{0} is within 256x256 size limit.".format(image_file.name))
            images.append((data, w, h, bpp))

        self.log.debug("Building icon...")
        header_section = []
        data_section = []

        # Set up the icon header
        icondir = ICONDIR()
        icondir.idReserved = 0
        icondir.idType = 1  # Icons are type 1
        icondir.idCount = len(images)
        header_section.append(icondir)
        
        icondir_size = ctypes.sizeof(icondir)
        iconentry_size = len(images) * ctypes.sizeof(ICONDIRENTRY())
        offset = icondir_size + iconentry_size  # calculate offset after entire header section
        for data, width, height, bpp in images:
            if width == 256:
                width = 0
            if height == 256:
                height = 0
            iconentry = ICONDIRENTRY()
            iconentry.bWidth = width
            iconentry.bHeight = height
            iconentry.bColorCount = 0           # Truecolor images have color count set to 0
            iconentry.bReserved = 0             # Reserved, must be 0
            iconentry.wPlanes = 1               # PNGs have 1 color plane
            iconentry.wBitCount = bpp           # bits per pixel from image_file
            iconentry.dwBytesInRes = len(data)  # Length of the image data
            # The data will be right after the icondir and iconentry
            iconentry.dwImageOffset = offset    # Offset of header and data so far
            offset += len(data)                 # Calculate new offset
            header_section.append(iconentry)
            data_section.append(data)

        self.log.debug("Writing icon...")
        with open(out_file, "wb") as f:
            # Write header sections followed by data sections
            for section in header_section + data_section:
                f.write(section)

    def run(self):
        """"""
        images = []
        if len(self.image_files) > 1:
            self.log.info("Assuming thumbnails have been provided in largest to smallest order, opening...")
            for image_file in self.image_files:
                with open(image_file,"rb") as f:
                    image = BytesIO(f.read())
                    image.name = image_file
                    image.seek(0)
                    images.append(image)
        else:
            self.log.info("Creating thumbnails...")
            for size in self.sizes:
                images.append(self._resize(self.image_files[0], size))
        self.log.info("Creating ICO file with {0} images...".format(len(images)))
        self._create_ico(self.out_file, images)
        #for image in images:
        #    image.close()


if __name__ == "__main__":
    print """
    PNG2ICO - Create a multisize ICO from a PNG
    Copyright (C) 2015  Oliver Morton (Sec-1 Ltd)    
    This program comes with ABSOLUTELY NO WARRANTY.
    This is free software, and you are welcome to redistribute it
    under certain conditions. See GPLv2 License.
""".format(__version__)

    def print_version():
        """Print command line version banner."""
        print """
.       .1111...          | Title: png2ico.py {0}    
    .10000000000011.   .. | Author: Oliver Morton (Sec-1 Ltd)
 .00              000...  | Email: oliverm-tools@sec-1.com
1                  01..   | Description:
                    ..    | Create a multisize ICO from a PNG file.
                   ..     |
GrimHacker        ..      |
                 ..       |
grimhacker.com  ..        |
@grimhacker    ..         |
----------------------------------------------------------------------------
""".format(__version__)

    class SizesAction(argparse.Action):
        def __call__(self, parser, namespace, values, option_string=None):
            # require values to be a list of integers between 1 and 256 (inclusive)
            max = 256
            min = 1
            for value in values:
                if value > max or value < min:
                    parser.error("sizes must be between {0} and {1}".format(max, min))
            setattr(namespace, self.dest, values)

    parser = argparse.ArgumentParser(description="Create a multisize ICO from a PNG")
    parser.add_argument("-V", "--version", help="Print version banner", action="store_true")
    parser.add_argument("-v", "--verbose", help="Debug logging", action="store_true")
    parser.add_argument("-i", "--input_file", help="Input PNG file", nargs="+", required=True)
    parser.add_argument("-o", "--output_file", help="Output ico filename", default="favicon.ico")
    parser.add_argument("-s", "--sizes", help="image sizes to create in ico", type=int, nargs="+", default=[256,48,32,16], action=SizesAction)
    args = parser.parse_args()

    if args.version:
        print_version()
        exit()

    if args.verbose:
        level = logging.DEBUG
    else:
        level = logging.INFO
    logging.basicConfig(level=level,
                        format="%(levelname)s: %(message)s")

    try:
        png2ico = PNG2ICO(image_files=args.input_file, out_file=args.output_file, sizes=args.sizes)
        png2ico.run()
    except Exception as e:
        logging.critical("Error running PNG2ICO: {0}".format(e))
    else:
        logging.info("Done. Check {0}".format(args.output_file))