50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import json
|
|
import socket
|
|
import ssl
|
|
|
|
from colony.settings import configuration
|
|
from colony.transport.framing import send_frame, recv_frame
|
|
|
|
|
|
class ConnectFailed(OSError):
|
|
pass
|
|
|
|
|
|
class Whisperer:
|
|
DEFAULT_TIMEOUT = 10
|
|
|
|
def __init__(self, host=None, port=None, cert_file=None, keyfile=None):
|
|
self.host = host or getattr(configuration.whisperer, "host", "127.0.0.1")
|
|
self.port = port or getattr(configuration.whisperer, "port", 9000)
|
|
self.cert_file = cert_file or getattr(configuration.whisperer, "cert_file", "certs/server.crt")
|
|
self.keyfile = keyfile or getattr(configuration.whisperer, "keyfile", "certs/server.key")
|
|
|
|
def _build_ssl_context(self):
|
|
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
|
context.minimum_version = ssl.TLSVersion.TLSv1_2
|
|
context.check_hostname = False
|
|
context.verify_mode = ssl.CERT_REQUIRED
|
|
context.load_verify_locations(cafile=self.cert_file)
|
|
context.load_cert_chain(certfile=self.cert_file, keyfile=self.keyfile)
|
|
return context
|
|
|
|
def send(self, payload, timeout=DEFAULT_TIMEOUT, connect_timeout=None):
|
|
data = json.dumps(payload).encode()
|
|
context = self._build_ssl_context()
|
|
|
|
try:
|
|
sock = socket.create_connection((self.host, self.port), timeout=connect_timeout or timeout)
|
|
except OSError as e:
|
|
raise ConnectFailed(f"Cannot connect to {self.host}:{self.port}: {e}") from e
|
|
|
|
with sock:
|
|
sock.settimeout(timeout)
|
|
try:
|
|
with context.wrap_socket(sock) as tls_sock:
|
|
send_frame(tls_sock, data)
|
|
response = recv_frame(tls_sock)
|
|
except socket.timeout:
|
|
raise TimeoutError("No response before timeout. ")
|
|
|
|
return json.loads(response.decode())
|