#!/usr/bin/python3
import sys
import hashlib
import random
import json
import pymysql


class keygen(object):

    def __init__(self, value):
        self.value = value

    #getkey
    def getkey(self, value):

        IsOrNot = None
        KeyRandom = str(random.getrandbits(16))
        IvRandom  = str(random.getrandbits(16))
        key = hashlib.md5(KeyRandom.encode('utf-8')).hexdigest()
        Iv = hashlib.md5(IvRandom.encode('utf-8')).hexdigest()

        db = pymysql.connect("localhost","root","Msfconsole123#","ksmPy" )
        # prepare a cursor object using cursor() method
        cursor = db.cursor()
        sql = "select * from KEYTABLE where id = %s" %str(value)

        cursor.execute(sql)
        results = cursor.fetchall()
        if not results:
            """print('Key/IV wasn\'t generated for this Id before')"""
            IsOrNot = False
        else:
            try:
                # Fetch all the rows in a list of lists.
                for row in results:
                    KMS_KEY = row[1]
                    KMS_IV = row[2]
                    # Now print fetched result
                    """print ("KEY = %s,IV = %s" % \
                    (KMS_KEY, KMS_IV))"""
            except:
                print ("Error: unable to fetch data")
            IsOrNot = True

        if(bool(IsOrNot) == False):

            sql = "INSERT INTO KEYTABLE(ID, \
              KSM_KEY, KSM_IV) \
              VALUES ('%s', '%s', '%s' )" % \
                (value, key, Iv)
            try:
                # Execute the SQL command
                cursor.execute(sql)
                # Commit your changes in the database
                db.commit()
                """print('New key & Iv successfully recorded ')
                print('ID : ' + value + ' Key: ' + key + ' Iv: ' + Iv)"""
            except:
                # Rollback in case there is any error
                db.rollback()
            # disconnect from server
            db.close()
        else:
           """print('key was already generated')"""

        return results