from flask import Flask, request, send_file
#import numpy
from PIL import Image, ImageChops
from io import BytesIO
import numpy

app = Flask(__name__)

usage = """
Post image Data to /desteg and receive a copy of that image with the least significant bits XORed with system entropy

Example: curl -H "Content-Type: application/png" --data-binary @image.png http://localhost:5002/desteg -o image_destegged.png

Updated 2022-08-11
"""

@app.route("/desteg", methods=["GET"])
def get_desteg():
    return usage

@app.route("/desteg", methods=["POST"])
def post_desteg():

    image_string = BytesIO(request.data)
    output_buffer = BytesIO()
    
    try:
        input_image = Image.open(image_string)
    except IOError:
        # insert relevant error return here
        pass



    # Create an array the same size
    # Stuff it full of entropy
    if (input_image.mode == "RGBA"):
        entropy_array = numpy.random.rand(
            input_image.size[1],
            input_image.size[0],
            4)*6
        entropy_image = Image.fromarray(
            entropy_array.astype('uint8'),
            mode="RGBA")
    elif (input_image.mode == "RGB"):
        entropy_array = numpy.random.rand(
            input_image.size[1],
            input_image.size[0],
            3)*6
        entropy_image = Image.fromarray(
            entropy_array.astype('uint8'), mode="RGB")
    else:
        # Insert relevant error return here
        pass

    # Add the entropy to input
    output_image = ImageChops.add(input_image, entropy_image)

    # save with same format as output
    output_image.save(output_buffer,format=input_image.format)

    #return buffer to start!
    output_buffer.seek(0)
    return send_file(output_buffer,
                     mimetype="image/" + input_image.format.lower())

@app.route("/")
def root():
    return usage

if __name__ == '__main__':
     app.run(port='5002',host="0.0.0.0")