#! /usr/bin/env python3

"""Reference implementation for Identity Bars.

This module contains functions to turn arbitrary bytes (for example hash
values and UUIDs) into identifiable pictures in PNG format.
"""

# Copyright (C) 2015-2026 David Heiko Kolf
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

_imgversion = {
    'version': 1,
    'modified': 'Sat, 14 Nov 2015 12:51:00 GMT'
}

import array
import base64
import email.utils
import hashlib
import re
import struct
import time
import zlib

class InvalidIdBarFilenameError(Exception):
    def __init__(self, filename):
        self.filename = filename

def _makechunk(chunktype, data):
    """Pack the data into a PNG chunk.

    Used internally.
    """
    length = struct.pack('!I', len(data))
    checksum = struct.pack('!I', zlib.crc32(data, zlib.crc32(chunktype)) & 0xffffffff)
    return length + chunktype + data + checksum

_pngsig = b'\x89PNG\r\n\x1a\n'
_iend = _makechunk(b'IEND', b'')

def _makeidhr(width, height, bitdepth, colortype, interlaced):
    """Create a PNG header.

    Used internally.
    """
    compression = 0
    filtermethod = 0
    return _makechunk(b'IHDR', struct.pack('!IIBBBBB', width, height,
        bitdepth, colortype, compression, filtermethod, interlaced))

def _scale3x(src, sx, sy, dst, dx, dy):
    """Interpolate a bitmap without introducing new colours.

    Used internally.

    src is a two dimensional array.

    sx and sy are pointing at the current pixel to be upscaled. There have
    to be neighbouring pixels to sx and sy.

    dst is a two dimensional array.
    dx and dy are pointing to the upper left pixel in the destination
    array. There have to be two further pixels to the right and to the
    bottom.

    The algorithm is described at http://www.scale2x.it/algorithm
    """
    a = src[sy-1][sx-1]
    b = src[sy-1][sx]
    c = src[sy-1][sx+1]
    d = src[sy][sx-1]
    e = src[sy][sx]
    f = src[sy][sx+1]
    g = src[sy+1][sx-1]
    h = src[sy+1][sx]
    i = src[sy+1][sx+1]

    x = b != h and d != f

    if x and d == b:
        dst[dy][dx] = d
    else:
        dst[dy][dx] = e
    if x and ((d == b and e != c) or (b == f and e != a)):
        dst[dy][dx+1] = b
    else:
        dst[dy][dx+1] = e
    if x and b == f:
        dst[dy][dx+2] = f
    else:
        dst[dy][dx+2] = e

    if x and ((d == h and a != e) or (b == d and e != g)):
        dst[dy+1][dx] = d
    else:
        dst[dy+1][dx] = e
    dst[dy+1][dx+1] = e
    if x and ((b == f and e != i) or (f == h and c != e)):
        dst[dy+1][dx+2] = f
    else:
        dst[dy+1][dx+2] = e

    if x and h == d:
        dst[dy+2][dx] = d
    else:
        dst[dy+2][dx] = e
    if x and ((f == h and e != g) or (h == d and e != i)):
        dst[dy+2][dx+1] = h
    else:
        dst[dy+2][dx+1] = e
    if x and h == f:
        dst[dy+2][dx+2] = f
    else:
        dst[dy+2][dx+2] = e

def _hsl2rgb(h, s, l):
    """Convert h, s, l values to RGB

    Used internally.
    """
    hn = h * 6.0 / 256.0
    c = (1 - abs(2 * l - 1)) * s
    m = l - 0.5 * c
    cm = int(255.0 * (c + m))
    xm = int(255.0 * (c * (1 - abs(hn % 2 - 1)) + m))
    zm = int(255.0 * m)
    if hn < 1:
        return (cm, xm, zm)
    elif hn < 2:
        return (xm, cm, zm)
    elif hn < 3:
        return (zm, cm, xm)
    elif hn < 4:
        return (zm, xm, cm)
    elif hn < 5:
        return (xm, zm, cm)
    else:
        return (cm, zm, xm)

def _makeicon(i1, i2):
    """Create a symmetrical icon using 15 bits provided by the two bytes
    i1 and i2.

    Returns an array with the placed bits, a border is added to simplify
    the scale3x implementation.

    Used internally.
    """
    def a(x):
        return 1&(i1>>x)
    def b(x):
        return 1&(i2>>x)
    return [
        [0,    0,    0,    0,    0,    0, 0],
        [0, a(0), a(5), b(2), a(5), a(0), 0],
        [0, a(1), a(6), b(3), a(6), a(1), 0],
        [0, a(2), a(7), b(4), a(7), a(2), 0],
        [0, a(3), b(0), b(5), b(0), a(3), 0],
        [0, a(4), b(1), b(6), b(1), a(4), 0],
        [0,    0,    0,    0,    0,    0, 0]
    ]

def _makesign(n):
    """Create a horizontal "signature" using all but the first 15 bits
    provided by the bytes in the array n.

    Returns an array with the placed bits, a border is added to simplify
    the scale3x implementation.

    Used internally.
    """
    sign = [[0]*(2*len(n)) for y in range(7)]
    sign[1][2] = 2 * (1&(n[1] >> 7)) # the icon didn't use this bit
    for i in range(0, len(n) - 2):
        b = n[i+2]
        for j in range(4):
            sign[j+2][i*2+1] = 2 * (1&(b >> j))
            sign[j+2][i*2+2] = 2 * (1&(b >> (j + 4)))
    return sign

def _makeimg(icon, sign, width, height):
    """Apply the scale3x algorithm to both icon and sign and return a
    two-dimensional array containing the combined image.

    Used internally.
    """
    img = [[0]*width for y in range(height)]
    for y in range(0, 5):
       for x in range(0, 5):
            _scale3x(icon, x+1, y+1, img, x*3, y*3)
    for y in range(0, 5):
        for x in range(0, len(sign[1])-4):
            _scale3x(sign, x+1, y+1, img, 16 + x*3, y*3)
    return img

def makepng(idhash):
    """Render an Identity Bar and return the PNG content as bytes based on
    the bytes given in idhash.
    """
    width = 4 + 6*len(idhash)
    height = 15

    iconcolor = _hsl2rgb(idhash[2], 1.0, 0.4)
    signcolor = _hsl2rgb(idhash[3], 0.2, 0.6)

    palette = array.array('B', [
          127, 255, 255,
          iconcolor[0], iconcolor[1], iconcolor[2],
          signcolor[0], signcolor[1], signcolor[2]
    ])

    icon = _makeicon(idhash[0], idhash[1])
    sign = _makesign(idhash)
    img = _makeimg(icon, sign, width, height)

    raw = array.array('B')
    for y in range(height):
      raw.append(0)
      for x in range(width):
        raw.append(img[y][x])

    ihdr = _makeidhr(width, height, 8, 3, 0)
    plte = _makechunk(b'PLTE', palette.tobytes())
    trns = _makechunk(b'tRNS', b'\0')
    idat = _makechunk(b'IDAT', zlib.compress(raw.tobytes()))
    png = _pngsig + ihdr + plte + trns + idat + _iend
    return png

def makefilename(prefix, idhash):
    """Create a filename starting with the given prefix and containing the
    bytes provided by idhash which can later be parsed by the function
    makefilecontent and wsgiapp.
    """
    return prefix + '-' + str(_imgversion['version']) + '-' + base64.urlsafe_b64encode(idhash).decode('ascii').rstrip('=') + '.png'

def makefilecontent(filename):
    """Parse a filename created by "makefilename" and return the PNG
    content as bytes.

    In case the filename cannot be parsed (by this version) the exception
    InvalidIdBarFilenameError is raised.
    """
    m = re.search(r'\-([0-9]+)\-([0-9A-Za-z\-\_]+)\.png$', filename)
    v = 0
    b64 = ''
    if m is not None:
        v = m.group(1)
        b64 = m.group(2)
    l = len(b64)
    p = 3 - ((l + 3) % 4)
    if int(v) != _imgversion['version'] or l < 6 or p >= 3:
        raise InvalidIdBarFilenameError(filename)
    return makepng(base64.urlsafe_b64decode(b64 + '=' * p))

def makedatauri(idhash):
    """Create a data-URI for the bytes given by idhash which can be
    directly included in an HTML file.
    """
    return ('data:image/png;base64,' +
        base64.b64encode(makepng(idhash)).decode('ascii'))

def wsgiapp(environ, start_response):
    """A WSGI implementation returning an appropriate PNG file for a
    path created by "makefilename" or a 404 response for invalid paths.
    """
    try:
        png = makefilecontent(environ.get('PATH_INFO', ''))
        start_response('200 OK', [
            ('Content-Type', 'image/png'),
            ('Content-Length', str(len(png))),
            ('Last-Modified', _imgversion['modified']),
            ('Expires', email.utils.formatdate(time.time() + 32e7, False, True))
        ])
        return (png,)
    except InvalidIdBarFilenameError:
        start_response('404 NOT FOUND', [('Content-Type', 'text/plain')])
        return ('Invalid ID bar URL',)

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description = 'Create ID bar PNG files')
    parser.add_argument('-x', '--hex')
    parser.add_argument('-f', '--fname')
    parser.add_argument('-o', '--output')
    args = parser.parse_args()
    fname = None
    content = None
    if args.hex is not None:
        idhash = bytes.fromhex(args.hex)
        fname = makefilename('idbar', idhash)
        content = makepng(idhash)
    if args.fname is not None:
        fname = args.fname
        content = makefilecontent(fname)
    if content is None:
        idhash = hashlib.sha256(str(time.time()).encode('ascii')).digest()
        print(idhash.hex())
        fname = makefilename('idbar', idhash)
        content = makepng(idhash)
    print(fname)
    output = fname
    if args.output is not None:
        output = args.output
    with open(output, 'wb') as f:
        f.write(content)