#!/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
# Some consts
PUBSUB_VERSION = '0.7.1'
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
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]
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))
self.server = ServerConfig(ip=server_ip, port=server_port, payload_limit=server_payload_limit)
# 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', [])
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 = {}
try:
self.acl = yaml.safe_load(open(args.acl))
except FileNotFoundError:
print(f"ACL configuration file {args.acl} 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),
}
# 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
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
subscriber = Subscriber(self, resp, request)
# Is there a basic auth in this request? If so, set up ACL
auth = request.headers.get('Authorization')
if auth:
subscriber.acl = await subscriber.parse_acl(auth)
# Subscribe the user before we deal with the potential backlog request and pings
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?
backlog = request.headers.get('X-Fetch-Since')
if backlog:
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])
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:
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.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()
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, connection: aiohttp.web.StreamResponse, request: aiohttp.web.BaseRequest):
self.connection = connection
self.acl = {}
self.server = server
self.lock = asyncio.Lock()
# 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).")
return acl
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
return 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}")
return {}
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
# Private payload?
if self.topics and self.topics[0] == 'private':
self.private = True
del self.topics[0] # Remove the private bit from topics now.
self.json['pubsub_timestamp'] = self.timestamp
self.json['pubsub_topics'] = self.topics
self.json['pubsub_path'] = path
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()