#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Thomas Roccia | IconDhash.py
# pip3 install lief
# pip3 install pillow
# resource: https://www.hackerfactor.com/blog/?/archives/529-Kind-of-Like-That.html

import lief
import os
import argparse
from PIL import Image


def extract_icon(exe):
    binary = lief.parse(exe)

    bin = binary.resources_manager
    ico = bin.icons
    ico = ico[0].save("peico.ico")
    return


def generate_icon_dhash(exe, hash_size = 8):
    extract_icon(exe)
    image = Image.open("peico.ico")
    image = image.convert('L').resize(
        (hash_size + 1, hash_size),
        Image.ANTIALIAS,
        )

    difference = []

    for row in range(hash_size):
        for col in range(hash_size):
            pixel_left = image.getpixel((col, row))
            pixel_right = image.getpixel((col + 1, row))
            difference.append(pixel_left > pixel_right)

    decimal_value = 0
    hex_string = []

    for index, value in enumerate(difference):
        if value:
            decimal_value += 2**(index % 8)

        if (index % 8) == 7:
            hex_string.append(hex(decimal_value)[2:].rjust(2, '0'))
            decimal_value = 0
    
    os.remove("peico.ico")

    return ''.join(hex_string)


# main function
def main():
    # select arguments
    parser = argparse.ArgumentParser(description='Generate icon dhash by Thomas Roccia')
    parser.add_argument("-f", "--file", help="Check domain list", required=True)
    args = parser.parse_args()

    # handle ctrl+c
    #signal.signal(signal.SIGINT, signal_handler)

    if args.file:
        try: 
            dhash = generate_icon_dhash(args.file)
            print("[+] dhash icon: %s" % dhash)
        except:
            print("[!] no icon available")


if __name__ == '__main__':
    main()