288 lines
10 KiB
Python
288 lines
10 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
Created on Sun Mar 22 10:44:50 2020
|
|
|
|
@author: cpan
|
|
"""
|
|
|
|
import sys
|
|
import selectors
|
|
import json
|
|
import io
|
|
import os
|
|
import struct
|
|
import hashlib
|
|
from Crypto import Random
|
|
from Crypto.PublicKey import RSA
|
|
|
|
|
|
request_search = {
|
|
"morpheus": "Follow the white rabbit. \U0001f430",
|
|
"ring": "In the caves beneath the Misty Mountains. \U0001f48d",
|
|
"\U0001f436": "\U0001f43e Playing ball! \U0001f3d0",
|
|
}
|
|
|
|
|
|
class Message:
|
|
def __init__(self, selector, sock, addr, passhash):
|
|
self.selector = selector
|
|
self.sock = sock
|
|
self.addr = addr
|
|
self._recv_buffer = b""
|
|
self._send_buffer = b""
|
|
self._jsonheader_len = None
|
|
self.jsonheader = None
|
|
self.request = None
|
|
self.response_created = False
|
|
# added for security
|
|
self.key = RSA.generate(2048, Random.new().read)
|
|
self.sessionkey = os.urandom(8)
|
|
self.password = passhash # pass phrase hash
|
|
self.clientkey = None # client public key
|
|
self.handshake = False
|
|
print(f"initiating message associated with {addr}")
|
|
|
|
def __del__(self):
|
|
print(f"destructing messange associlated with {self.addr}")
|
|
|
|
def _set_selector_events_mask(self, mode):
|
|
"""Set selector to listen for events: mode is 'r', 'w', or 'rw'."""
|
|
if mode == "r":
|
|
events = selectors.EVENT_READ
|
|
elif mode == "w":
|
|
events = selectors.EVENT_WRITE
|
|
elif mode == "rw":
|
|
events = selectors.EVENT_READ | selectors.EVENT_WRITE
|
|
else:
|
|
raise ValueError(f"Invalid events mask mode {repr(mode)}.")
|
|
self.selector.modify(self.sock, events, data=self)
|
|
|
|
def _read(self):
|
|
try:
|
|
# Should be ready to read
|
|
data = self.sock.recv(4096)
|
|
except BlockingIOError:
|
|
# Resource temporarily unavailable (errno EWOULDBLOCK)
|
|
pass
|
|
else:
|
|
if data:
|
|
self._recv_buffer += data
|
|
else:
|
|
raise RuntimeError("Peer closed.")
|
|
|
|
def _write(self):
|
|
if self._send_buffer:
|
|
print("sending", repr(self._send_buffer), "to", self.addr)
|
|
try:
|
|
# Should be ready to write
|
|
sent = self.sock.send(self._send_buffer)
|
|
except BlockingIOError:
|
|
# Resource temporarily unavailable (errno EWOULDBLOCK)
|
|
pass
|
|
else:
|
|
self._send_buffer = self._send_buffer[sent:]
|
|
if not self.handshake:
|
|
self._set_selector_events_mask("r")
|
|
# Close when the buffer is drained. The response has been sent.
|
|
elif sent and not self._send_buffer:
|
|
self.close()
|
|
|
|
def _json_encode(self, obj, encoding):
|
|
return json.dumps(obj, ensure_ascii=False).encode(encoding)
|
|
|
|
def _json_decode(self, json_bytes, encoding):
|
|
tiow = io.TextIOWrapper(
|
|
io.BytesIO(json_bytes), encoding=encoding, newline=""
|
|
)
|
|
obj = json.load(tiow)
|
|
tiow.close()
|
|
return obj
|
|
|
|
def _create_message(
|
|
self, *, 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 = self._json_encode(jsonheader, "utf-8")
|
|
message_hdr = struct.pack(">H", len(jsonheader_bytes))
|
|
message = message_hdr + jsonheader_bytes + content_bytes
|
|
return message
|
|
|
|
def _create_response_json_content(self):
|
|
action = self.request.get("action")
|
|
if action == "search":
|
|
query = self.request.get("value")
|
|
answer = request_search.get(query) or f'No match for "{query}".'
|
|
content = {"result": answer}
|
|
else:
|
|
content = {"result": f'Error: invalid action "{action}".'}
|
|
content_encoding = "utf-8"
|
|
response = {
|
|
"content_bytes": self._json_encode(content, content_encoding),
|
|
"content_type": "text/json",
|
|
"content_encoding": content_encoding,
|
|
}
|
|
return response
|
|
|
|
def _create_response_binary_content(self):
|
|
response = {
|
|
"content_bytes": b"First 10 bytes of request: "
|
|
+ self.request[:10],
|
|
"content_type": "binary/custom-server-binary-type",
|
|
"content_encoding": "binary",
|
|
}
|
|
return response
|
|
|
|
def _create_handshake_content(self, data):
|
|
response = {
|
|
"content_bytes": data,
|
|
"content_type": "binary/custom-server-binary-type",
|
|
"content_encoding": "binary",
|
|
}
|
|
return response
|
|
|
|
def process_events(self, mask):
|
|
if mask & selectors.EVENT_READ:
|
|
self.read()
|
|
if mask & selectors.EVENT_WRITE:
|
|
self.write()
|
|
|
|
def read(self):
|
|
self._read()
|
|
|
|
if self._jsonheader_len is None:
|
|
self.process_protoheader()
|
|
|
|
if self._jsonheader_len is not None:
|
|
if self.jsonheader is None:
|
|
self.process_jsonheader()
|
|
|
|
if self.jsonheader:
|
|
if self.request is None:
|
|
self.process_request()
|
|
|
|
if not self.handshake:
|
|
self._jsonheader_len = None
|
|
self.jsonheader = None
|
|
|
|
def write(self):
|
|
if self.request:
|
|
if not self.response_created:
|
|
self.create_response()
|
|
|
|
self._write()
|
|
|
|
def close(self):
|
|
print(f"closing connection to {self.addr}\n")
|
|
try:
|
|
self.selector.unregister(self.sock)
|
|
except Exception as e:
|
|
print(
|
|
f"error: selector.unregister() exception for",
|
|
f"{self.addr}: {repr(e)}",
|
|
)
|
|
|
|
try:
|
|
self.sock.close()
|
|
except OSError as e:
|
|
print(
|
|
f"error: socket.close() exception for",
|
|
f"{self.addr}: {repr(e)}",
|
|
)
|
|
finally:
|
|
# Delete reference to socket object for garbage collection
|
|
self.sock = None
|
|
|
|
def process_protoheader(self):
|
|
hdrlen = 2
|
|
if len(self._recv_buffer) >= hdrlen:
|
|
self._jsonheader_len = struct.unpack(
|
|
">H", self._recv_buffer[:hdrlen]
|
|
)[0]
|
|
self._recv_buffer = self._recv_buffer[hdrlen:]
|
|
|
|
def process_jsonheader(self):
|
|
hdrlen = self._jsonheader_len
|
|
if len(self._recv_buffer) >= hdrlen:
|
|
self.jsonheader = self._json_decode(
|
|
self._recv_buffer[:hdrlen], "utf-8"
|
|
)
|
|
self._recv_buffer = self._recv_buffer[hdrlen:]
|
|
for reqhdr in (
|
|
"byteorder",
|
|
"content-length",
|
|
"content-type",
|
|
"content-encoding",
|
|
):
|
|
if reqhdr not in self.jsonheader:
|
|
raise ValueError(f'Missing required header "{reqhdr}".')
|
|
|
|
def send_my_secrets(self):
|
|
halfsesskey = self.sessionkey
|
|
encryptedsess = self.clientkey.encrypt(halfsesskey, None)[0]
|
|
sesskeyhash = hashlib.sha3_512(halfsesskey).hexdigest().encode()
|
|
mypubkey = self.key.publickey().exportKey('PEM')
|
|
mypubkeyhash = hashlib.sha3_512(mypubkey).hexdigest().encode()
|
|
mysecret = mypubkey + mypubkeyhash + encryptedsess + sesskeyhash
|
|
response = self._create_handshake_content(mysecret)
|
|
message = self._create_message(**response)
|
|
self._send_buffer += message
|
|
print(f'Sending server public key, hash, session key, hash...')
|
|
|
|
def process_request(self):
|
|
content_len = self.jsonheader["content-length"]
|
|
if not len(self._recv_buffer) >= content_len:
|
|
return
|
|
data = self._recv_buffer[:content_len]
|
|
self._recv_buffer = self._recv_buffer[content_len:]
|
|
if self.clientkey is None: # data is client public key plus hash
|
|
key = data[:-128]
|
|
keyhash = data[-128:]
|
|
if hashlib.sha3_512(key).hexdigest().encode() == keyhash:
|
|
self.clientkey = RSA.importKey(key)
|
|
else:
|
|
self.sock.close()
|
|
raise Exception('Client plublic key does not match key hash')
|
|
self.send_my_secrets()
|
|
self._set_selector_events_mask("w")
|
|
elif not self.handshake: # data is passhass + sesshash
|
|
passphrase = self.key.decrypt(bytes(data))
|
|
sesshash = passphrase[-128:]
|
|
passhash = passphrase[:-128]
|
|
if passhash != self.password:
|
|
self.sock.close()
|
|
raise Exception('Password mismatch')
|
|
if hashlib.sha3_512(self.sessionkey).hexdigest().encode() != sesshash:
|
|
self.sock.close()
|
|
raise Exception('Session key hash mismatch')
|
|
self.handshake = True
|
|
print('Handshake Done.')
|
|
self._set_selector_events_mask("r")
|
|
else:
|
|
if self.jsonheader["content-type"] == "text/json":
|
|
encoding = self.jsonheader["content-encoding"]
|
|
self.request = self._json_decode(data, encoding)
|
|
print("received request", repr(self.request), "from", self.addr)
|
|
else:
|
|
# Binary or unknown content-type
|
|
self.request = data
|
|
print(
|
|
f'received {self.jsonheader["content-type"]} request from',
|
|
self.addr,
|
|
)
|
|
# Set selector to listen for write events, we're done reading.
|
|
self._set_selector_events_mask("w")
|
|
|
|
def create_response(self):
|
|
if self.jsonheader["content-type"] == "text/json":
|
|
response = self._create_response_json_content()
|
|
else:
|
|
# Binary or unknown content-type
|
|
response = self._create_response_binary_content()
|
|
message = self._create_message(**response)
|
|
self.response_created = True
|
|
self._send_buffer += message |