##  Copyright (c) 2011-2023, Art Obrezan.
##  All rights reserved. See the license at the end of the file.

import hashlib
import socket
import struct

def test():
    with socket.create_connection(('127.0.0.1', 3306)) as connection:
        packet = read_packet(connection)
        if not is_error_packet(packet):
            scramble = handshake_scramble(packet)
            packet = auth_packet(scramble, 'username', 'password', 'mysql')
            write_packet(connection, packet, 1)
            packet = read_packet(connection)
            # query
            res = doquery(connection, 'show tables')
            return res


##-----------------------------------------------------------------------------
## HANDSHAKE PARSING
##-----------------------------------------------------------------------------

def handshake_protocol(packet):
    return (packet[0])

def handshake_scramble(packet):
    pos = packet.find(0)
    scramble = bytearray(20)
    for i in range(8):
        scramble[i] = packet[pos+5+i]
    for i in range(12):
        scramble[i+8] = packet[pos+32+i]
    return bytes(scramble)

def handshake_version(packet):
    pos = packet.find(0)
    return str(packet[1:pos], encoding='latin1')

def handshake_threadid(packet):
    pos = packet.find(0)
    return (packet[pos+1] +
            packet[pos+2]*0xFF +
            packet[pos+3]*0xFFFF +
            packet[pos+4]*0xFFFFFF)

def handshake_capabilities(packet):
    pos = packet.find(0)
    return (packet[pos+14] +
            packet[pos+15]*0xFF)

def handshake_language(packet):
    pos = packet.find(0)
    return (packet[pos+16])

def handshake_status(packet):
    pos = packet.find(0)
    return (packet[pos+17] +
            packet[pos+18]*0xFF)


##-----------------------------------------------------------------------------

TOKEN_LEN = 20

MAX_PACKET_SIZE = 1024 * 1024  # in bytes
UTF8_GENERAL_CI = 33


def auth_packet(scramble, user, password=None, database=None):
    buf_user = bytes(user, encoding='latin1') + b'\0'

    buf_database = b''
    if database:
        buf_database = bytes(database, encoding='latin1') + b'\0'

    buf_auth = b'\0'
    if password:
        token = password_to_token(password, scramble)
        buf_auth = bytes(bytearray([TOKEN_LEN]) + token)

    flags = client_capabilities()
    if database:
        flags = flags + 0x08  # client-connect-with-db

    buf_flags = struct.pack('<I', flags)
    buf_packet_size = struct.pack('<I', MAX_PACKET_SIZE)
    buf_charset = struct.pack('<I', UTF8_GENERAL_CI)

    return (buf_flags + buf_packet_size + buf_charset + bytearray(20) +
            buf_user + buf_auth + buf_database)


def password_to_token(password, scramble):
    stage1 = hashlib.sha1(bytes(password, encoding='latin1')).digest()
    stage2 = hashlib.sha1(stage1).digest()
    digest = hashlib.sha1(scramble + stage2).digest()
    token = bytearray(TOKEN_LEN)
    for i in range(TOKEN_LEN):
        token[i] = digest[i] ^ stage1[i]
    return bytes(token)

def client_capabilities():
    """ http://dev.mysql.com/doc/internals/en/capability-flags.html """
    return (0x0200 + # client-protocol-41
            0x8000 + # client-secure-connection (authentication::native41)
            0x0100 + # client-ignore-space
            0x2000)  # client-transactions


##-----------------------------------------------------------------------------
## QUERY
##-----------------------------------------------------------------------------

def doquery(connection, query):
    send_query(connection, query)
    packet = read_packet(connection)
    if is_error_packet(packet):
        return 'error'
    elif is_ok_packet(packet):
        return 'ok'
    else:
        return parse_data_packets(connection, packet)

def parse_data_packets(connection, packet):
    fields_number, l = decode_length_coded_binary(packet, 0)
    # read mysql field packets
    for i in range(fields_number):
        packet = read_packet(connection)
    # read eof packet
    packet = read_packet(connection)
    # read row packets data
    res = []
    while True:
        packet = read_packet(connection)
        if is_eof_packet(packet):
            break
        row = parse_row_packet(packet)
        res.append(row)
    return res

def parse_row_packet(packet):
    row = []
    packet_length = len(packet)
    pos = 0
    while True:
        field_length, field_pos = decode_length_coded_binary(packet, pos)
        if field_length == -1: # column value == NULL
            row.append("NULL")
            pos = pos + field_pos
        else:
            start = pos + field_pos
            end = start + field_length
            value = str(packet[start:end], encoding='utf-8')
            # value = packet[start:end]
            row.append(value)
            pos = end
        if pos >= packet_length:
            break
    return row

def nth_field_entry(n, buf):
    pos = 0
    len = 0
    offset = 0
    for i in n:
        pos = pos + offset + len
        len, offset = decode_length_coded_binary(buf, pos)
    pos = pos + offset
    return (pos, len)

def send_query(connection, query):
    COM_QUERY = b'\3'
    data = COM_QUERY + bytes(query, encoding='utf8')
    packet_number = 0
    write_packet(connection, data, packet_number)


##-----------------------------------------------------------------------------
## LENGTH CODED BINARY
##-----------------------------------------------------------------------------

def decode_length_coded_binary(buf, pos):
    len = None
    val = buf[pos]
    if 251 > val:
        len = (val, 1)
    elif 251 == val:
        len = (-1, 1)
    elif 252 == val:
        int16 = decode_int16(buf, pos)
        len = (int16, 3)
    elif 253 == val:
        int24 = decode_int24(buf, pos)
        len = (int24, 4)
    elif 254 == val:
        int64 = decode_int64(buf, pos)
        len = (int64, 9)
    return len

def decode_int16(buf, pos):
    data = (buf[pos+1] +
            buf[pos+2]*0xFF)
    return data

def decode_int24(buf, pos):
    data = (buf[pos+1] +
            buf[pos+2]*0xFF +
            buf[pos+3]*0xFFFF)
    return data

def decode_int64(buf, pos):
    data = (buf[pos+1] +
            buf[pos+2]*0xFF +
            buf[pos+3]*0xFFFF +
            buf[pos+4]*0xFFFFFF +
            buf[pos+5]*0XFFFFFFFF +
            buf[pos+6]*0xFFFFFFFFFF +
            buf[pos+7]*0XFFFFFFFFFFFF +
            buf[pos+8]*0XFFFFFFFFFFFFFF)
    return data


##-----------------------------------------------------------------------------
## PACKET READING/WRITING/CHECKING
##-----------------------------------------------------------------------------

def read_packet(connection):
    len, num = read_packet_header(connection)
    buf = connection.recv(len)
    return buf

def read_packet_header(connection):
    buf = connection.recv(4)
    len = buf[0] + buf[1]*0xFF + buf[2]*0xFFFF
    num = buf[3]
    return len, num

def write_packet(connection, data, packet_number):
    write_packet_header(connection, data, packet_number)
    connection.sendall(data)

def write_packet_header(connection, data, packet_number):
    len4bytes = struct.pack('<I', len(data))
    len3bytes = len4bytes[:3]
    buf = len3bytes + bytearray([packet_number & 0xFF])
    connection.sendall(buf)

def is_error_packet(packet):
    return (0xFF == packet[0])

def is_ok_packet(packet):
    return (0x00 == packet[0])

def is_eof_packet(packet):
    return ((0xFE == packet[0]) and (5 == len(packet)))



## ----------------------------------------------------------------------------
##
##  Copyright (c) 2011-2023, Art Obrezan
##  All rights reserved.
##
##  Redistribution and use in source and binary forms, with or without
##  modification, are permitted provided that the following conditions
##  are met:
##  1. Redistributions of source code must retain the above copyright
##     notice, this list of conditions and the following disclaimer.
##  2. Redistributions in binary form must reproduce the above copyright
##     notice, this list of conditions and the following disclaimer in the
##     documentation and/or other materials provided with the distribution.
##  3. Use in source and binary forms are not permitted in projects
##     under GNU General Public Licenses and its derivatives.
##
##  THIS SOFTWARE IS PROVIDED BY ART OBREZAN ''AS IS'' AND ANY EXPRESS OR
##  IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
##  OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
##  IN NO EVENT SHALL ART OBREZAN BE LIABLE FOR ANY  DIRECT, INDIRECT,
##  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
##  NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;  LOSS OF USE,
##  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
##  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
##  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
##  THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
##
## ----------------------------------------------------------------------------