#!/usr/bin/env python3 # Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """PyPubSub - a simple publisher/subscriber service written in Python 3""" import asyncio import aiohttp.web import aiofile import os import time import json import yaml import netaddr import binascii import base64 import argparse import collections import plugins.ldap import plugins.sqs import typing import signal import uuid # Some consts PUBSUB_VERSION = '0.7.4' PUBSUB_CONTENT_TYPE = 'application/vnd.pypubsub-stream' PUBSUB_DEFAULT_PORT = 2069 PUBSUB_DEFAULT_IP = '0.0.0.0' PUBSUB_DEFAULT_MAX_PAYLOAD_SIZE = 102400 PUBSUB_DEFAULT_BACKLOG_SIZE = 0 PUBSUB_DEFAULT_BACKLOG_AGE = 0 PUBSUB_BAD_REQUEST = "I could not understand your request, sorry! Please see https://pubsub.apache.org/api.html \ for usage documentation.\n" PUBSUB_PAYLOAD_RECEIVED = "Payload received, thank you very much!\n" PUBSUB_NOT_ALLOWED = "You are not authorized to deliver payloads!\n" PUBSUB_BAD_PAYLOAD = "Bad payload type. Payloads must be JSON dictionary objects, {..}!\n" PUBSUB_PAYLOAD_TOO_LARGE = "Payload is too large for me to serve, please make it shorter.\n" PUBSUB_WRITE_TIMEOUT = 0.35 # If we can't deliver to a pipe within N seconds, drop it. class ServerConfig(typing.NamedTuple): ip: str port: int payload_limit: int tls_port: int tls_ctx: typing.Any class BacklogConfig(typing.NamedTuple): max_age: int queue_size: int storage: typing.Optional[str] class Configuration: server: ServerConfig backlog: BacklogConfig payloaders: typing.List[netaddr.ip.IPNetwork] oldschoolers: typing.List[str] secure_topics: typing.Optional[typing.List[str]] def __init__(self, yml: dict): # LDAP Settings self.ldap = None lyml = yml.get('clients', {}).get('ldap') if isinstance(lyml, dict): self.ldap = plugins.ldap.LDAPConnection(lyml) # SQS? self.sqs = yml.get('sqs') # Main server config server_ip = yml['server'].get('bind', PUBSUB_DEFAULT_IP) server_port = int(yml['server'].get('port', PUBSUB_DEFAULT_PORT)) server_payload_limit = int(yml['server'].get('max_payload_size', PUBSUB_DEFAULT_MAX_PAYLOAD_SIZE)) tls_port = 0 tls_ctx = None # TLS support, if configured if 'tls' in yml['server'] and isinstance(yml['server']['tls'], dict): for required_element in ("port", "cert", "key", ): assert yml['server']['tls'].get(required_element), f"TLS: configuration option '{required_element}' is missing or invalid, cannot enable TLS!" import ssl tls_port = int(yml['server']['tls']['port']) # Create TLS context and load cert+key tls_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH) assert os.path.isfile(yml['server']['tls']['cert']), f"Could not locate domain certificate file {yml['server']['tls']['cert']}" assert os.path.isfile(yml['server']['tls']['key']), f"Could not locate domain certificate key {yml['server']['tls']['key']}" tls_ctx.load_cert_chain(yml['server']['tls']['cert'], yml['server']['tls']['key']) # Add intermediate cert chain, if available if 'chain' in yml['server']['tls']: assert os.path.isfile(yml['server']['tls']['chain']), f"Could not locate domain certificate chain {yml['server']['tls']['chain']}" tls_ctx.load_verify_locations(yml['server']['tls']['chain']) self.server = ServerConfig(ip=server_ip, port=server_port, payload_limit=server_payload_limit, tls_port=tls_port, tls_ctx=tls_ctx) # Backlog settings bma = yml['server'].get('backlog', {}).get('max_age', PUBSUB_DEFAULT_BACKLOG_AGE) if isinstance(bma, str): bma = bma.lower() if bma.endswith('s'): bma = int(bma.replace('s', '')) elif bma.endswith('m'): bma = int(bma.replace('m', '')) * 60 elif bma.endswith('h'): bma = int(bma.replace('h', '')) * 3600 elif bma.endswith('d'): bma = int(bma.replace('d', '')) * 86400 bqs = yml['server'].get('backlog', {}).get('size', PUBSUB_DEFAULT_BACKLOG_SIZE) bst = yml['server'].get('backlog', {}).get('storage') self.backlog = BacklogConfig(max_age=bma, queue_size=bqs, storage=bst) # Payloaders - clients that can post payloads self.payloaders = [netaddr.IPNetwork(x) for x in yml['clients'].get('payloaders', [])] # Binary backwards compatibility self.oldschoolers = yml['clients'].get('oldschoolers', []) # Secure topics, if any self.secure_topics = set(yml['clients'].get('secure_topics', []) or []) class Server: """Main server class, responsible for handling requests and publishing events """ yaml: dict config: Configuration subscribers: list pending_events: asyncio.Queue backlog: list last_ping = typing.Type[float] server: aiohttp.web.Server def __init__(self, args: argparse.Namespace): self.yaml = yaml.safe_load(open(args.config)) self.config = Configuration(self.yaml) self.subscribers = [] self.pending_events = asyncio.Queue() self.backlog = [] self.last_ping = time.time() self.acl_file = args.acl self.acl = {} self.load_acl() def load_acl(self): """Loads ACL from file""" try: self.acl = yaml.safe_load(open(self.acl_file)) print(f"Loaded ACL from {self.acl_file}") except FileNotFoundError: print(f"ACL configuration file {self.acl_file} not found, private events will not be broadcast.") async def poll(self): """Polls for new stuff to publish, and if found, publishes to whomever wants it.""" while True: payload: Payload = await self.pending_events.get() bad_subs: list = await payload.publish(self.subscribers) self.pending_events.task_done() # Cull subscribers we couldn't deliver payload to. for bad_sub in bad_subs: print("Culling %r due to connection errors" % bad_sub) try: self.subscribers.remove(bad_sub) except ValueError: # Already removed elsewhere pass async def handle_request(self, request: aiohttp.web.BaseRequest): """Generic handler for all incoming HTTP requests""" resp: typing.Union[aiohttp.web.Response, aiohttp.web.StreamResponse] # Define response headers first... headers = { 'Server': 'PyPubSub/%s' % PUBSUB_VERSION, 'X-Subscribers': str(len(self.subscribers)), 'X-Requests': str(self.server.requests_count), } subscriber = Subscriber(self, request) # Is there a basic auth in this request? If so, set up ACL auth = request.headers.get('Authorization') if auth: await subscriber.parse_acl(auth) # Are we handling a publisher payload request? (PUT/POST) if request.method in ['PUT', 'POST']: ip = netaddr.IPAddress(request.remote) allowed = False for network in self.config.payloaders: if ip in network: allowed = True break # Check for secure topics payload_topics = set(request.path.split("/")) if any(x in self.config.secure_topics for x in payload_topics): allowed = False # Figure out which secure topics we need permission for: which_secure = [x for x in self.config.secure_topics if x in payload_topics] # Is the user allowed to post to all of these secure topics? if subscriber.secure_topics and all(x in subscriber.secure_topics for x in which_secure): allowed = True if not allowed: resp = aiohttp.web.Response(headers=headers, status=403, text=PUBSUB_NOT_ALLOWED) return resp if request.can_read_body: try: if request.content_length and request.content_length > self.config.server.payload_limit: resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_PAYLOAD_TOO_LARGE) return resp body = await request.text() as_json = json.loads(body) assert isinstance(as_json, dict) # Payload MUST be an dictionary object, {...} pl = Payload(request.path, as_json) self.pending_events.put_nowait(pl) # Add to backlog? if self.config.backlog.queue_size > 0: self.backlog.append(pl) # If backlog has grown too large, delete the first (oldest) item in it. while len(self.backlog) > self.config.backlog.queue_size: del self.backlog[0] resp = aiohttp.web.Response(headers=headers, status=202, text=PUBSUB_PAYLOAD_RECEIVED) return resp except json.decoder.JSONDecodeError: resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_REQUEST) return resp except AssertionError: resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_PAYLOAD) return resp # Is this a subscriber request? (GET) elif request.method == 'GET': resp = aiohttp.web.StreamResponse(headers=headers) # We do not support HTTP 1.0 here... if request.version.major == 1 and request.version.minor == 0: return resp # Subscribe the user before we deal with the potential backlog request and pings subscriber.connection = resp self.subscribers.append(subscriber) resp.content_type = PUBSUB_CONTENT_TYPE try: resp.enable_chunked_encoding() await resp.prepare(request) # Is the client requesting a backlog of items? epoch_based_backlog = request.headers.get('X-Fetch-Since') cursor_based_backlog = request.headers.get('X-Fetch-Since-Cursor') if epoch_based_backlog: # epoch-based backlog search try: backlog_ts = int(backlog) except ValueError: # Default to 0 if we can't parse the epoch backlog_ts = 0 # If max age is specified, force the TS to minimum that age if self.config.backlog.max_age > 0: backlog_ts = max(backlog_ts, int(time.time() - self.config.backlog.max_age)) # For each item, publish to client if new enough. for item in self.backlog: if item.timestamp >= backlog_ts: await item.publish([subscriber]) if cursor_based_backlog and len(cursor_based_backlog) == 36: # UUID4 cursor-based backlog search # For each item, publish to client if it was published after this cursor is_after_cursor = False for item in self.backlog: if item.cursor == cursor_based_backlog: # Found cursor, mark it! is_after_cursor = True elif is_after_cursor: # This is after the cursor, stream it await item.publish([subscriber]) while True: await subscriber.ping() if subscriber not in self.subscribers: # If we got dislodged somehow, end session break await asyncio.sleep(5) # We may get exception types we don't have imported, so grab ANY exception and kick out the subscriber except Exception: # don't catch fatal exceptions pass if subscriber in self.subscribers: self.subscribers.remove(subscriber) return resp elif request.method == 'HEAD': resp = aiohttp.web.Response(headers=headers, status=204, text="") return resp # I don't know this type of request :/ (DELETE, PATCH, etc) else: resp = aiohttp.web.Response(headers=headers, status=400, text=PUBSUB_BAD_REQUEST) return resp async def write_backlog_storage(self): previous_backlog = [] while True: if self.config.backlog.storage: try: backlog_list = self.backlog.copy() if backlog_list != previous_backlog: previous_backlog = backlog_list async with aiofile.AIOFile(self.config.backlog.storage, 'w+') as afp: offset = 0 for item in backlog_list: js =json.dumps({ 'timestamp': item.timestamp, 'topics': item.topics, 'json': item.json, 'private': item.private }) + '\n' await afp.write(js, offset=offset) offset += len(js) await afp.fsync() except Exception as e: print(f"Could not write to backlog file {self.config.backlog.storage}: {e}") await asyncio.sleep(10) def read_backlog_storage(self): if self.config.backlog.storage and os.path.exists(self.config.backlog.storage): try: readlines = 0 with open(self.config.backlog.storage, 'r') as fp: for line in fp.readlines(): js = json.loads(line) readlines += 1 ppath = "/".join(js['topics']) if js['private']: ppath = '/private/' + ppath payload = Payload(ppath, js['json'], js['timestamp']) self.backlog.append(payload) if self.config.backlog.queue_size < len(self.backlog): self.backlog.pop(0) except Exception as e: print(f"Error while reading backlog: {e}") print(f"Read {readlines} objects from {self.config.backlog.storage}, applied {len(self.backlog)} to backlog.") async def server_loop(self, loop: asyncio.BaseEventLoop): self.server = aiohttp.web.Server(self.handle_request) runner = aiohttp.web.ServerRunner(self.server) await runner.setup() site = aiohttp.web.TCPSite(runner, self.config.server.ip, self.config.server.port) await site.start() print("==== PyPubSub v/%s starting... ====" % PUBSUB_VERSION) print("==== Serving up PubSub goodness at %s:%s ====" % ( self.config.server.ip, self.config.server.port)) if self.config.server.tls_ctx: site_tls = aiohttp.web.TCPSite(runner, self.config.server.ip, self.config.server.tls_port, ssl_context=self.config.server.tls_ctx) await site_tls.start() print("==== Serving up PubSub TLS goodness at %s:%s ====" % ( self.config.server.ip, self.config.server.tls_port)) if self.config.sqs: for key, config in self.config.sqs.items(): loop.create_task(plugins.sqs.get_payloads(self, config)) self.read_backlog_storage() loop.create_task(self.write_backlog_storage()) await self.poll() def run(self): loop = asyncio.get_event_loop() # add a signal handler for SIGUSR2 to reload the ACL from disk try: loop.add_signal_handler(signal.SIGUSR2, self.load_acl) except ValueError: pass try: loop.run_until_complete(self.server_loop(loop)) except KeyboardInterrupt: pass loop.close() class Subscriber: """Basic subscriber (client) class. Holds information about the connection and ACL""" acl: dict topics: typing.List[typing.List[str]] def __init__(self, server: Server, request: aiohttp.web.BaseRequest): self.connection: typing.Optional[aiohttp.web.StreamResponse] = None self.acl = {} self.server = server self.lock = asyncio.Lock() self.secure_topics = [] # Set topics subscribed to self.topics = [] for topic_batch in request.path.split(','): sub_to = [x for x in topic_batch.split('/') if x] self.topics.append(sub_to) # Is the client old and expecting zero-terminators? self.old_school = False for ua in self.server.config.oldschoolers: if ua in request.headers.get('User-Agent', ''): self.old_school = True break async def parse_acl(self, basic: str): """Sets the ACL if possible, based on Basic Auth""" try: decoded = str(base64.decodebytes(bytes(basic.replace('Basic ', ''), 'ascii')), 'utf-8') u, p = decoded.split(':', 1) if u in self.server.acl: acl_pass = self.server.acl[u].get('password') if acl_pass and acl_pass == p: acl = self.server.acl[u].get('acl', {}) # Vet ACL for user assert isinstance(acl, dict), f"ACL for user {u} " \ f"must be a dictionary of sub-IDs and topics, but is not." # Make sure each ACL segment is a list of topics for k, v in acl.items(): assert isinstance(v, list), f"ACL segment {k} for user {u} is not a list of topics!" print(f"Client {u} successfully authenticated (and ACL is valid).") self.acl = acl self.secure_topics = set(self.server.acl[u].get('topics', []) or []) elif self.server.config.ldap: acl = {} groups = await self.server.config.ldap.get_groups(u,p) # Make sure each ACL segment is a list of topics for k, v in self.server.config.ldap.acl.items(): if k in groups: assert isinstance(v, dict), f"ACL segment {k} for user {u} is not a dictionary of segments!" for segment, topics in v.items(): print(f"Enabling ACL segment {segment} for user {u}") assert isinstance(topics, list), f"ACL segment {segment} for user {u} is not a list of topics!" acl[segment] = topics self.acl = acl except binascii.Error as e: pass # Bad Basic Auth params, bail quietly except AssertionError as e: print(e) print(f"ACL configuration error: ACL scheme for {u} contains errors, setting ACL to nothing.") except Exception as e: print(f"Basic unknown exception occurred: {e}") async def ping(self): """Generic ping-back to the client""" js = b"%s\n" % json.dumps({"stillalive": time.time()}).encode('utf-8') if self.old_school: js += b"\0" async with self.lock: await asyncio.wait_for(self.connection.write(js), timeout=PUBSUB_WRITE_TIMEOUT) class Payload: """A payload (event) object sent by a registered publisher.""" def __init__(self, path: str, data: dict, timestamp: typing.Optional[float] = None): self.json = data self.timestamp = timestamp or time.time() self.topics = [x for x in path.split('/') if x] self.private = False self.cursor = str(uuid.uuid4()) # Event cursor for playback - UUID4 style # Private payload? if self.topics and self.topics[0] == 'private': self.private = True del self.topics[0] # Remove the private bit from topics now. # Set standard pubsub meta data in the payload self.json['pubsub_timestamp'] = self.timestamp self.json['pubsub_topics'] = self.topics self.json['pubsub_path'] = path self.json['pubsub_cursor'] = self.cursor async def publish(self, subscribers: typing.List[Subscriber]): """Publishes an object to all subscribers using those topics (or a sub-set thereof)""" js = b"%s\n" % json.dumps(self.json).encode('utf-8') ojs = js + b"\0" bad_subs = [] for sub in subscribers: # If a private payload, check ACL and bail if not a match if self.private: can_see = False for key, private_topics in sub.acl.items(): if all(el in self.topics for el in private_topics): can_see = True break if not can_see: continue # If subscribed to all the topics, tell a subscriber about this for topic_batch in sub.topics: if all(el in self.topics for el in topic_batch): try: if sub.old_school: async with sub.lock: await asyncio.wait_for(sub.connection.write(ojs), timeout=PUBSUB_WRITE_TIMEOUT) else: async with sub.lock: await asyncio.wait_for(sub.connection.write(js), timeout=PUBSUB_WRITE_TIMEOUT) except Exception: bad_subs.append(sub) break return bad_subs if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("--config", help="Configuration file to load (default: pypubsub.yaml)", default="pypubsub.yaml") parser.add_argument("--acl", help="ACL Configuration file to load (default: pypubsub_acl.yaml)", default="pypubsub_acl.yaml") cliargs = parser.parse_args() pubsub_server = Server(cliargs) pubsub_server.run()