# -*- 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(bytes(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 RSA.importKey(pubkey), halfsesskey def main(): HDRLEN = 2 _recv_buffer = bytearray(b"") if len(sys.argv) != 5: print("usage:", sys.argv[0], " ") 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(None) # 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) # acknowledge session key and passphrase # TODO send encrypted key and passphrase back 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()