from starlette.endpoints import HTTPEndpoint
from starlette.responses import JSONResponse
from database import database
import sqlite3
import os

class Exploit(HTTPEndpoint):
    async def get(self, request):
        pass
    async def put(self, request):
        # Get file from form
        form = await request.form()
        filename = form['file'].filename
        contents = await form['file'].read()
        # Write file to exploits folder
        open("./exploits/" + filename, 'w+b').write(contents)
        os.chmod("./exploit/" + filename, "+x")

        # Get exploit data
        service_id = request.query_params['service_id']

        # Insert into database
        try:
            database.insert_exploit(filename, service_id)
        except sqlite3.IntegrityError as e:
            return JSONResponse({"error":str(e)})

        # Get and return result
        ret = database.get_exploits(filename)[0] # Get first and only result
        return JSONResponse({"filename":ret[1], "service":ret[4]})