import socket
import threading
import hashlib
import rsa
import sqlite3

ID = 'id'
NAME = 'name'
VOTES = 'votes'
TABLE_VOTE_DICT = \
    {
        ID: 0,
        NAME: 1,
        VOTES: 2
    }
CREATE_VOTE_TABLE = 'CREATE TABLE IF NOT EXISTS elections (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, name varchar, votes INTEGER)'
INC_VOTE = "UPDATE elections SET votes = '%s' WHERE name = '%s'"
VOTE_DICT = \
    {
        0: 'Putin',
        1: 'Zyuganov',
        2: 'Zhirinovsky',
        3: 'Mironov'
    }

def _incVote_(c, vote, name):
    c.execute(INC_VOTE%(vote+1,name))

def _checkIfTableExists_(c):
    c.execute(CREATE_VOTE_TABLE)

def handleVote(vote):
    vote &=  ((1 << 20) - 1)
    print(vote)
    conn = sqlite3.connect('elections.sqlite')
    c = conn.cursor()
    c.execute('SELECT * FROM elections')
    row = c.fetchone()
    while row is not None:
        if row[TABLE_VOTE_DICT[NAME]] == VOTE_DICT[vote]:
            _incVote_(c, row[TABLE_VOTE_DICT[VOTES]], VOTE_DICT[vote])
            break
        row = c.fetchone()
    conn.commit()
    c.close()
    conn.close()

def addCandidate(name):
    conn = sqlite3.connect('elections.sqlite')
    c = conn.cursor()
    c.execute("INSERT INTO elections (name,votes) VALUES ('%s','%i')"%(name,0))
    conn.commit()
    c.close()
    conn.close()

def resetElection():
    conn = sqlite3.connect('elections.sqlite')
    c = conn.cursor()
    c.execute("UPDATE elections SET votes = 0")
    conn.commit()
    c.close()
    conn.close()

LOGIN = 'login'
PASSWORD = 'password'
VOTED = 'voted'
ERR_PASSWORD_TOO_SHORT = 'Password is too short!'
ERR_LOGIN_EXISTS = 'Login already exists!'
MIN_PASS = 6
MAX_PASS = 20
TABLE_DICT = \
    {
        ID: 0,
        LOGIN: 1,
        PASSWORD: 2,
        VOTED: 3
    }
CREATE_TABLE = 'CREATE TABLE IF NOT EXISTS '\
               +'users (id INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, '\
               +'login varchar, password varchar, voted INTEGER)'
ADD_USER = "INSERT INTO users (login, password, voted) VALUES ('%s','%s', 0)"
MARK_AS_VOTED = "UPDATE users SET voted = 1 WHERE login = '%s'"
NULL_VOTES = "UPDATE users SET voted = 0"

def _checkPassLength_(password):
    length = len(password)
    return length >= MIN_PASS or length <= MAX_PASS

def _checkLoginUniqueness_(c, login):
    # TODO: simplify the query
    c.execute("SELECT * FROM users")
    for row in c:
        if (login == row[TABLE_DICT[LOGIN]]):
            return False
    return True

def _str2hashstr_(string):
    tNumber = string
    hashObj = hashlib.sha256()
    byteArray = str.encode(string)
    hashObj.update(byteArray)
    return hashObj.hexdigest()

def _initDB_():
    conn = sqlite3.connect('users.sqlite')
    c = conn.cursor()
    c.execute(CREATE_TABLE)
    return (conn, c)

def _closeDB_(conn, c):
    conn.commit()
    c.close()
    conn.close()

def createAccount(login, password):
    if not _checkPassLength_(password):
        print(ERR_PASSWORD_TOO_SHORT)
        return False
    conn, c = _initDB_()
    if not _checkLoginUniqueness_(c, login):
        print(ERR_LOGIN_EXISTS)
        return False
    hash = _str2hashstr_(password)
    c.execute(ADD_USER%(login,hash))
    _closeDB_(conn, c)
    return True

def authentification(login, password):
    conn, c = _initDB_()
    # TODO: simplify the query
    c.execute("SELECT * FROM users")
    for row in c:
        if login == row[TABLE_DICT[LOGIN]]:
            _closeDB_(conn, c)
            return _str2hashstr_(password) == row[TABLE_DICT[PASSWORD]]
    #print('Login not found!')
    _closeDB_(conn, c)

def markAsVoted(login):
    ret = False
    conn, c = _initDB_()
    # TODO: simplify the query
    c.execute("SELECT * FROM users")
    for row in c:
        if row[TABLE_DICT[LOGIN]] == login:
            c.execute(MARK_AS_VOTED%login)
            ret = True
    _closeDB_(conn, c)
    return ret

def checkIfVoted(login):
    conn, c = _initDB_()
    # TODO: simplify the query
    c.execute("SELECT * FROM users")
    for row in c:
        if row[TABLE_DICT[LOGIN]] == login:
            return bool(row[TABLE_DICT[VOTED]])
    _closeDB_(conn, c)

def nullVotes():
    conn, c = _initDB_()
    c.execute("SELECT * FROM users")
    c.execute(NULL_VOTES)
    _closeDB_(conn, c)

def powMod (n ,pow ,mod):
    res=1
    while pow>0:
        if (pow&1):
            res = (res * n) % mod
        n = (n * n) % mod
        pow>>=1
    return res

def number2hashnum(number):
    tNumber = number
    hashObj = hashlib.sha256()
    byteNumber = bytearray()
    while tNumber > 0:
        byteNumber.append(tNumber & 0xFF)
        tNumber >>= 8
    byteNumber.reverse()
    hashObj.update(byteNumber)
    hash = hashObj.digest()
    hashArray = bytearray(hash)
    hashNumber = 0
    hashArray.reverse()
    length = len(hashArray)
    i = 0
    while i<length :
        hashNumber += hashArray[i] << (i * 8)
        i += 1
    return hashNumber

def checkVote(R,S,D,N):
    res = powMod(S,D,N)
    hash = number2hashnum(R)
    return res==hash

def sendKeys(client,N,D):
    client.send((str(N)+' '+str(D)+' ').encode())

    print('Keys sent')

def sendBulletin(client):
    bulletin = 'President elections 2018\nChoose the candidate:\n1.Putin\n2.Zuganov\n3.Zhirinovsky\n4.Mironov\n'
    client.send((bulletin).encode())

def signVote(h_, c, n):
    return powMod(h_, c, n)

def receiveVote():
    voteSock = socket.socket()
    voteSock.bind(('', 9091))
    voteSock.listen(5)
    while(True):
        clientSock, clientAddr = voteSock.accept()
        vote = clientSock.recv(2048)
        R, S = parseMessageTo2Ints(vote)
        isValid = checkVote(R,S,D,N)
        print(isValid)
        if(isValid):
            handleVote(R)
        clientSock.close()
        print('Vote received')

def userThread(clientSock):
    while(True):
        message = clientSock.recv(2048)
        message = message.decode().split(' ')
        if(message[0]=="reg"):
            clientSock.send(('1').encode()) if createAccount(message[1],message[2]) else clientSock.send(('0').encode())
        if(message[0]=="log" and authentification(message[1],message[2])):
            if(checkIfVoted(message[1])):
                print('User already voted')
                clientSock.send(('2').encode())
            else:
                markAsVoted(message[1])
                clientSock.send(('1').encode())
                print('Ready to send')
                break
        if(message[0]=="log" and not authentification(message[1],message[2])):
            print('Login failed')
            clientSock.send(('0').encode())
    sendKeys(clientSock, N, D)
    clientSock.recv(2048)
    sendBulletin(clientSock)
    h_ = clientSock.recv(2048)
    signedVote = signVote(int(h_.decode()), C, N)
    clientSock.send(str(signedVote).encode())
    clientSock.close()

def parseMessageTo2Ints(message):
    message = message.decode()
    message = message.split(' ')
    return int(message[0]), int(message[1])

def ParametrPrivkeys():
    (pubkey, privkey) = rsa.newkeys(1024)
    D = privkey.e
    C = privkey.d
    N = privkey.n
    Q = privkey.q
    P = privkey.p
    return D, C, N, Q, P

D, C, N, Q, P = ParametrPrivkeys()
servSock = socket.socket()
servSock.bind(('', 9090))
servSock.listen(5)
resetElection()
nullVotes()
print('Server launched!')

voteThread = threading.Thread(target=receiveVote)
voteThread.start()

userThreads = list()

while(True):
    clientSock, clientAddr = servSock.accept()
    userThreads.append(threading.Thread(target=userThread,args=[clientSock]))
    userThreads[len(userThreads)-1].start()

print("I'm off")