wiki: Only match usernames when a word starts with @

[ipfire.org.git] / src / backend / ratelimit.py

#!/usr/bin/python

import datetime

from . import misc

class RateLimiter(misc.Object):
	def handle_request(self, request, handler, minutes, limit):
		return RateLimiterRequest(self.backend, request, handler,
			minutes=minutes, limit=limit)


class RateLimiterRequest(misc.Object):
	def init(self, request, handler, minutes, limit):
		self.request = request
		self.handler = handler

		# Save the limits
		self.minutes = minutes
		self.limit   = limit

		# What is the current time?
		self.now = datetime.datetime.utcnow()

		# When to expire?
		self.expires_at = self.now + datetime.timedelta(minutes=self.minutes + 1)

		self.prefix = "-".join((
			self.__class__.__name__,
			self.request.host,
			self.request.path,
			self.request.method,
			self.request.remote_ip,
		))

	async def is_ratelimited(self):
		"""
			Returns True if the request is prohibited by the rate limiter
		"""
		counter = await self.get_counter()

		# The client is rate-limited when more requests have been
		# received than allowed.
		if counter >= self.limit:
			return True

		# Increment the counter
		await self.increment_counter()

		# If not ratelimited, write some headers
		self.write_headers(counter=counter)

	@property
	def key(self):
		return "%s-%s" % (self.prefix, self.now.strftime("%Y-%m-%d-%H:%M"))

	@property
	def keys_to_check(self):
		for minute in range(self.minutes + 1):
			when = self.now - datetime.timedelta(minutes=minute)

			yield "%s-%s" % (self.prefix, when.strftime("%Y-%m-%d-%H:%M"))

	async def get_counter(self):
		"""
			Returns the number of requests that have been done in
			recent time.
		"""
		async with await self.backend.cache.pipeline() as p:
			for key in self.keys_to_check:
				await p.get(key)

			# Run the pipeline
			res = await p.execute()

		# Return the sum
		return sum((int(e) for e in res if e))

	def write_headers(self, counter):
		# Send the limit to the user
		self.handler.set_header("X-Rate-Limit-Limit", self.limit)

		# Send the user how many requests are left for this time window
		self.handler.set_header("X-Rate-Limit-Remaining", self.limit - counter)

		# Send when the limit resets
		self.handler.set_header("X-Rate-Limit-Reset", self.expires_at.strftime("%s"))

	async def increment_counter(self):
		async with await self.backend.cache.pipeline() as p:
			# Increment the key
			await p.incr(self.key)

			# Set expiry
			await p.expireat(self.key, self.expires_at)

			# Run the pipeline
			await p.execute()