nopaste: Implement PROXY protocol on the TCP service

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

#!/usr/bin/python3

import asyncio
import datetime
import io
import ipaddress
import logging
import magic
import struct
import tornado.iostream
import tornado.tcpserver

from . import base
from .misc import Object
from .decorators import *

# Setup logging
log = logging.getLogger(__name__)

CHUNK_SIZE = 1024 ** 2

class Nopaste(Object):
	def _get_paste(self, query, *args, **kwargs):
		return self.db.fetch_one(Paste, query, *args, **kwargs)

	def create(self, content, subject=None, mimetype=None, expires=None, account=None, address=None):
		# Store the blob
		blob_id = self._store_blob(content)

		# Guess the mimetype if none set
		if not mimetype:
			mimetype = magic.from_buffer(content, mime=True)

		if expires:
			expires = datetime.datetime.utcnow() + datetime.timedelta(seconds=expires)

		# http://blog.00null.net/easily-generating-random-strings-in-postgresql/
		paste = self._get_paste("""
			INSERT INTO
				nopaste
			(
				uuid,
				subject,
				content,
				time_expires,
				address,
				mimetype,
				size,
				blob_id
			)
			VALUES
			(
				random_slug(), %s, %s, %s, %s, %s, %s, %s
			)
			RETURNING
				*
			""", subject, content, expires or None, address, mimetype, len(content), blob_id,
		)

		# Log result
		log.info("Created a new paste (%s) of %s byte(s) from %s (%s - %s)" % (
			paste.uuid, paste.size, paste.address, paste.asn or "N/A", paste.country or "N/A",
		))

		return paste

	def _fetch_blob(self, id):
		blob = self.db.get("""
			SELECT
				data
			FROM
				nopaste_blobs
			WHERE
				id = %s
			""", id,
		)

		if blob:
			return blob.data

	def _store_blob(self, data):
		"""
			Stores the blob by sending it to the database and returning its ID
		"""
		blob = self.db.get("""
			INSERT INTO
				nopaste_blobs
			(
				data
			)
			VALUES
			(
				%s
			)
			ON CONFLICT
			(
				digest(data, 'sha256')
			)
			DO UPDATE SET
				last_uploaded_at = CURRENT_TIMESTAMP
			RETURNING
				id
			""", data,
		)

		# Return the ID
		return blob.id

	def get(self, uuid):
		paste = self._get_paste("""
			SELECT
				*
			FROM
				nopaste
			WHERE
				uuid = %s
			AND (
				expires_at >= CURRENT_TIMESTAMP
			OR
				expires_at IS NULL
			)
			""", uuid,
		)

		return paste

	def get_content(self, uuid):
		res = self.db.get("SELECT content FROM nopaste \
			WHERE uuid = %s", uuid)

		if res:
			return bytes(res.content)

	def cleanup(self):
		"""
			Removes all expired pastes and removes any unneeded blobs
		"""
		# Remove all expired pastes
		self.db.execute("""
			DELETE FROM
				nopaste
			WHERE
				expires_at < CURRENT_TIMESTAMP
		""")

		# Remove unneeded blobs
		self.db.execute("""
			DELETE FROM
				nopaste_blobs
			WHERE NOT EXISTS
			(
				SELECT
					1
				FROM
					nopaste
				WHERE
					nopaste.blob_id = nopaste_blobs.id
			)
		""")


class Paste(Object):
	def init(self, id, data):
		self.id, self.data = id, data

	# UUID

	@property
	def uuid(self):
		return self.data.uuid

	# Subject

	@property
	def subject(self):
		return self.data.subject

	# Created At

	@property
	def created_at(self):
		return self.data.created_at

	time_created = created_at

	# Expires At

	@property
	def expires_at(self):
		return self.data.expires_at

	time_expires = expires_at

	# Blob

	@lazy_property
	def blob(self):
		return self.backend.nopaste._fetch_blob(self.data.blob_id)

	content = blob

	# Size

	@property
	def size(self):
		return self.data.size

	# MIME Type

	@property
	def mimetype(self):
		return self.data.mimetype or "application/octet-stream"

	# Address

	@property
	def address(self):
		return self.data.address

	# Location

	@lazy_property
	def location(self):
		 return self.backend.location.lookup("%s" % self.address)

	# ASN

	@lazy_property
	def asn(self):
		if self.location and self.location.asn:
			return self.backend.location.get_asn(self.location.asn)

	# Country

	@lazy_property
	def country(self):
		if self.location and self.location.country_code:
			return self.backend.location.get_country(self.location.country_code)

	# Viewed?

	def viewed(self):
		"""
			Call this when this paste has been viewed/downloaded/etc.
		"""
		self.db.execute("""
			UPDATE
				nopaste
			SET
				last_accessed_at = CURRENT_TIMESTAMP,
				views = views + 1
			WHERE
				id = %s
			""", self.id,
		)


# PROXY Protocol Implementation
PROXY_CMD_LOCAL     = 0
PROXY_CMD_PROXY     = 1

PROXY_FAMILY_UNSPEC = 0
PROXY_FAMILY_INET   = 1
PROXY_FAMILY_INET6  = 2

PROXY_PROTO_UNSPEC  = 0
PROXY_PROTO_STREAM  = 1
PROXY_PROTO_DGRAM   = 2

class ProxyError(Exception):
	pass


class ProxyUnsupportedError(ProxyError):
	pass


class Service(tornado.tcpserver.TCPServer):
	def __init__(self, config, use_proxy=True, **kwargs):
		# Initialize backend
		self.backend = base.Backend(config)

		# Expect PROXY headers?
		self.use_proxy = use_proxy

		super().__init__(**kwargs)

	async def handle_stream(self, stream, address):
		buffer = io.BytesIO()

		# Parse the PROXY header
		if self.use_proxy:
			try:
				address = await self._parse_proxyv2_header(stream, address)

			# Close the stream on any proxy errors
			except ProxyError as e:
				log.error("Proxy Error: %s" % e)

				return stream.close()

			# If we received no result, this connection is not supposed to be relayed
			if not address:
				return

		# Read the entire stream
		try:
			while True:
				chunk = await stream.read_bytes(CHUNK_SIZE, partial=True)

				log.debug("Read a chunk of %s byte(s)" % len(chunk))

				# Write the chunk into the buffer
				buffer.write(chunk)

				# If we have read less then we have reached the end
				if len(chunk) < CHUNK_SIZE:
					break

		# End if the stream closed unexpectedly
		except tornado.iostream.StreamClosedError as e:
			return

		log.debug("Finished reading payload")

		# Process address
		address, port = address

		# Store this into the database
		with self.backend.db.transaction():
			paste = self.backend.nopaste.create(
				buffer.getvalue(),
				subject="Streamed Upload",
				address=address,
			)

		# Format a response message
		message = "https://nopaste.ipfire.org/view/%s\n" % paste.uuid

		# Send the message
		await stream.write(message.encode("utf-8"))

		# We are done, close the stream
		stream.close()

	async def _parse_proxyv2_header(self, stream, address):
		"""
			Parses the PROXYv2 header and returns the real client's IP address
		"""
		src_address, src_port, dst_address, dst_port = None, None, None, None

		log.debug("Parsing PROXY connection from %s:%s" % address)

		# Header
		proxy_hdr_v2  = struct.Struct("!12BBBH")
		proxy_addr_v6 = struct.Struct("!16B16BHH")
		proxy_addr_v4 = struct.Struct("!IIHH")

		# Try to read the header into a buffer
		buffer = await stream.read_bytes(proxy_hdr_v2.size)

		if len(buffer) < proxy_hdr_v2.size:
			raise ProxyError("Header too short")

		# Parse the header
		*signature, ver_cmd, fam_prot, length = proxy_hdr_v2.unpack(buffer)

		# Check signature
		if not signature == [13, 10, 13, 10, 0, 13, 10, 81, 85, 73, 84, 10]:
			raise ProxyError("Incorrect signature")

		# Extract version and command
		version = (ver_cmd >> 4)
		command = (ver_cmd & 0x0f)

		# Check protocol version
		if not version == 2:
			raise ProxyError("Incorrect protocol version")

		# Handle LOCAL commands
		if command == PROXY_CMD_LOCAL:
			pass

		# Handle PROXY commands
		elif command == PROXY_CMD_PROXY:
			pass # Fallthrough

		# We don't know any other commands
		else:
			log.debug("Unknown PROXY command %02x" % command)
			return

		# Extract protocol
		family   = (fam_prot >> 4)
		protocol = (fam_prot & 0x0f)

		# LOCAL command use family == AF_UNSPEC
		if command == PROXY_CMD_LOCAL and family == PROXY_FAMILY_UNSPEC:
			pass

		# We accept IPv6 and IPv4
		elif family in (PROXY_FAMILY_INET6, PROXY_FAMILY_INET):
			pass

		# Everything else we don't know how to handle here
		else:
			raise ProxyError("Unknown family %s" % family)

		# LOCAL commands use protocol == UNSPEC
		if command == PROXY_CMD_LOCAL and protocol == PROXY_PROTO_UNSPEC:
			pass

		# Otherwise we only support TCP
		elif protocol == PROXY_PROTO_STREAM:
			pass

		# We don't know how to handle anything else here
		else:
			raise ProxyUnsupportedError("Unknown or unsupported protocol %s" % protocol)

		# Read the next part of the header into the buffer
		buffer = await stream.read_bytes(length)

		# Check if we read enough data
		if len(buffer) < length:
			raise ProxyError("Header too short")

		# Unpack IPv6 addresses
		if family == PROXY_FAMILY_INET6:
			addresses = proxy_addr_v6.unpack(buffer[:proxy_addr_v6.size])

			src_address = ipaddress.IPv6Address(bytes(addresses[:16]))
			src_port    = addresses[-2]
			dst_address = ipaddress.IPv6Address(bytes(addresses[16:32]))
			dst_port    = addresses[-1]

			# Truncate buffer
			buffer = buffer[proxy_addr_v6.size:]

		# Unpack IPv4 addresses
		elif family == PROXY_FAMILY_INET:
			src_address, dst_address, src_port, dst_port = proxy_addr_v4.unpack(
				buffer[:proxy_addr_v4.size],
			)

			# Convert IP addresses
			src_address = ipaddress.IPv4Address(src_address)
			dst_address = ipaddress.IPv4Address(dst_address)

			# Truncate buffer
			buffer = buffer[proxy_addr_v4.size:]

		# Handle UNSPEC
		elif family == PROXY_FAMILY_UNSPEC:
			pass

		# Log the result
		if src_address and dst_address:
			log.debug("Accepted new connection from %s:%s to %s:%s" \
				% (src_address, src_port, dst_address, dst_port))

		# Return the source address
		return "%s" % src_address, src_port