Socket_Programming/app_secure_client.py

# -*- coding: utf-8 -*-
"""
Created on Sun Mar 22 10:46:20 2020

@author: cpan

Simple client implementation with timeout feature to prevent hanging

"""

import sys
import socket
import json
import io
import struct
import hashlib
from Crypto import Random
from Crypto.PublicKey import RSA

def create_request(action, value):
    if action == "search":
        return dict(
            type="text/json",
            encoding="utf-8",
            content=dict(action=action, value=value),
        )
    else:
        return dict(
            type="binary/custom-client-binary-type",
            encoding="binary",
            content=bytes(action + value, encoding="utf-8"),
        )


def _json_encode(obj, encoding):
    return json.dumps(obj, ensure_ascii=False).encode(encoding)


def _json_decode(json_bytes, encoding):
    tiow = io.TextIOWrapper(
        io.BytesIO(json_bytes), encoding=encoding, newline=""
    )
    obj = json.load(tiow)
    tiow.close()
    return obj


def _create_message(*, content_bytes, content_type, content_encoding):
    jsonheader = {
        "byteorder": sys.byteorder,
        "content-type": content_type,
        "content-encoding": content_encoding,
        "content-length": len(content_bytes),
    }
    jsonheader_bytes = _json_encode(jsonheader, "utf-8")
    message_hdr = struct.pack(">H", len(jsonheader_bytes))
    message = message_hdr + jsonheader_bytes + content_bytes
    return message


def create_message(request):
    content = request["content"]
    content_type = request["type"]
    content_encoding = request["encoding"]
    if content_type == "text/json":
        req = {
            "content_bytes": _json_encode(content, content_encoding),
            "content_type": content_type,
            "content_encoding": content_encoding,
        }
    else:
        req = {
            "content_bytes": content,
            "content_type": content_type,
            "content_encoding": content_encoding,
        }
    message = _create_message(**req)
    return message


def read_to_buffer(sock, length, buf : bytearray):
    while len(buf) < length:
        data = sock.recv(4096)
        buf += data


def process_response(sock, addr, jsonheader, _recv_buffer):
    content_len = jsonheader["content-length"]
    read_to_buffer(sock, content_len, _recv_buffer)
    data = _recv_buffer[:content_len]
    _recv_buffer = _recv_buffer[content_len:]
    if jsonheader["content-type"] == "text/json":
        encoding = jsonheader["content-encoding"]
        response = _json_decode(data, encoding)
        print("received response", repr(response), "from", addr)
        content = response
        result = content.get("result")
        print(f"got result: {result}")
    else:
        # Binary or unknown content-type
        response = data
        print(
            f'received {jsonheader["content-type"]} response from',
            addr,
        )
        content = response
        print(f"got response: {repr(content)}")


def send_client_key(sock, data):
        response = {
            "content_bytes": data,
            "content_type": "binary/custom-server-binary-type",
            "content_encoding": "binary",
        }
        message = _create_message(**response)
        print(f'Sending client public key and hash...')
        sock.sendall(message)

def recv_server_secrets(sock, _recv_buffer, key):
    HDRLEN = 2
    # read and process protoheader (fixed length of 2 byte in network order)
    read_to_buffer(sock, HDRLEN, _recv_buffer)
    _jsonheader_len = struct.unpack(">H", _recv_buffer[:HDRLEN])[0]
    _recv_buffer = _recv_buffer[HDRLEN:]

    # read and process jsonheader
    read_to_buffer(sock, _jsonheader_len, _recv_buffer)
    jsonheader = _json_decode(_recv_buffer[:_jsonheader_len], "utf-8")
    _recv_buffer = _recv_buffer[_jsonheader_len:]
    for reqhdr in ("byteorder", "content-length", "content-type",
                   "content-encoding",):
        if reqhdr not in jsonheader:
            raise ValueError(f'Missing required header "{reqhdr}".')

    content_len = jsonheader["content-length"]
    read_to_buffer(sock, content_len, _recv_buffer)
    data = _recv_buffer[:content_len]
    _recv_buffer = _recv_buffer[content_len:]
    sesskeyhash = data[-128:]
    data = data[:-128]
    encryptedsess = data[-256:]
    # halfsesskey = key.decrypt(encryptedsess)
    data = data[:-256]
    pubkeyhash = data[-128:]
    pubkey = data[:-128]
    if hashlib.sha3_512(pubkey).hexdigest().encode() != pubkeyhash:
        sock.shutdown(socket.SHUT_RDWR)
        sock.close()
        raise Exception('Server public key does not match hash')
    if hashlib.sha3_512(halfsesskey).hexdigest().encode() != sesskeyhash:
        sock.shutdown(socket.SHUT_RDWR)
        sock.close()
        raise Exception('Server public key does not match hash')
    print(f"Received server secrets correctly.")
    return pubkey, halfsesskey


def main():
    HDRLEN = 2
    _recv_buffer = bytearray(b"")

    if len(sys.argv) != 5:
        print("usage:", sys.argv[0], "<host> <port> <action> <value>")
        sys.exit(1)

    host, port = sys.argv[1], int(sys.argv[2])
    addr = (host, port)
    print("starting connection to", addr)
    sock = socket.create_connection(addr)
    sock.settimeout(10)

    # security handshake
    key = RSA.generate(2048, Random.new().read)
    publickey = key.publickey().exportKey('PEM')
    keyhash = hashlib.sha3_512(publickey).hexdigest().encode()
    data = publickey + keyhash
    send_client_key(sock, data)
    # receiving server secrets - public key and session key
    serverkey, sessionkey = recv_server_secrets(sock, _recv_buffer, key)

    action, value = sys.argv[3], sys.argv[4]
    request = create_request(action, value)
    message = create_message(request)
    print(f"sending request {repr(message)} to {addr}")
    sock.sendall(message)

    # read and process protoheader (fixed length of 2 byte in network order)
    read_to_buffer(sock, HDRLEN, _recv_buffer)
    _jsonheader_len = struct.unpack(">H", _recv_buffer[:HDRLEN])[0]
    _recv_buffer = _recv_buffer[HDRLEN:]

    # read and process jsonheader
    read_to_buffer(sock, _jsonheader_len, _recv_buffer)
    jsonheader = _json_decode(_recv_buffer[:_jsonheader_len], "utf-8")
    _recv_buffer = _recv_buffer[_jsonheader_len:]
    for reqhdr in ("byteorder", "content-length", "content-type",
                   "content-encoding",):
        if reqhdr not in jsonheader:
            raise ValueError(f'Missing required header "{reqhdr}".')

    # read and process response
    process_response(sock, addr, jsonheader, _recv_buffer)

    sock.shutdown(socket.SHUT_RDWR)
    sock.close()


if __name__ == '__main__':
    main()