nopaste: Tidy up the cleanup task

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

#!/usr/bin/python3

import asyncio
import datetime
import io
import logging
import magic
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)

		uid = None
		if account:
			uid = account.uid

		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,
				uid,
				mimetype,
				size,
				blob_id
			)
			VALUES
			(
				random_slug(), %s, %s, %s, %s, %s, %s, %s, %s
			)
			RETURNING
				*
			""", subject, content, expires or None, address, uid, 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
			(
				data
			)
			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 _update_lastseen(self, uuid):
		self.db.execute("UPDATE nopaste SET time_lastseen = NOW(), views = views + 1 \
			WHERE uuid = %s", uuid)

	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)

	# Legacy

	@property
	def account(self):
		return None


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

		super().__init__(**kwargs)

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

		# 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():
			uuid = self.backend.nopaste.create(
				buffer.getvalue(),
				subject="Streamed Upload",
				address=address,
			)

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

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

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