From 3c088027c045b60e00a799d86601d9ccbe0e9682 Mon Sep 17 00:00:00 2001
From: iou1name <iou1name@national.shitposting.agency>
Date: Sun, 6 Oct 2019 17:01:46 -0400
Subject: [PATCH] upload from url

---
 README.md            |  2 +-
 saddle.py            | 67 ++++++++++++++++++++++++++++++++++++++++++--
 templates/index.html |  1 +
 3 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 1496de8..f9d92b2 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@ A file hosting service similar to Pomf and Uguu but without the public nature.
 ## Requirements  
 Python 3.7+  
 PostgreSQL 11.5+  
-Python packages: `wheel gunicorn aiohttp aiohttp_jinja2 asyncpg uvloop`  
+Python packages: `wheel gunicorn aiohttp aiohttp_jinja2 asyncpg uvloop requests`  
 
 ## Install  
 ```
diff --git a/saddle.py b/saddle.py
index ebaf99e..1037948 100644
--- a/saddle.py
+++ b/saddle.py
@@ -3,6 +3,7 @@
 A file hosting service similar to Pomf and Uguu but without the public nature.
 """
 import os
+import time
 import string
 import random
 import datetime
@@ -13,6 +14,7 @@ import aiohttp_jinja2
 from aiohttp_jinja2 import render_template
 import asyncpg
 import uvloop
+import requests
 
 import config
 import buckler_aiohttp
@@ -29,9 +31,15 @@ async def index(request):
 		return render_template("index.html", request, locals())
 	data = await request.post()
 	rand_name = bool(data.get('rand_name'))
+
 	files = []
 	for filefield in data.getall('files'):
+		if not filefield:
+			continue
 		files.append(handle_filefield(filefield, rand_name=rand_name))
+	if data.get('url'):
+		files.append(handle_url(data.get('url'), rand_name=rand_name))
+
 	if data.get('delete_this'):
 		delete_num = data.get('delete_num', '')
 		delete_type = data.get('delete_type', '')
@@ -74,8 +82,31 @@ def handle_filefield(filefield, rand_name=True):
 	with open(os.path.join(config.upload_dir, filename), 'wb') as file:
 		file.write(filefield.file.read())
 
-	t = (prefix, filename)
-	return t
+	tup = (prefix, filename)
+	return tup
+
+
+def handle_url(url, rand_name=True):
+	"""Handles a posted URL."""
+	try:
+		filename, data = download_file(url)
+	except ValueError:
+		return None
+
+	filename = safe_filename(filename)
+	if not filename:
+		rand_name = True
+	prefix = get_rand_chars()
+	if rand_name:
+		filename = prefix + os.path.splitext(filename)[1]
+	else:
+		filename = prefix + '_' + filename
+
+	with open(os.path.join(config.upload_dir, filename), 'wb') as file:
+		file.write(data)
+
+	tup = (prefix, filename)
+	return tup
 
 
 def safe_filename(filename=''):
@@ -95,6 +126,36 @@ def get_rand_chars(n=8):
 	return "".join(chars)
 
 
+def download_file(url, timeout=10, max_file_size=config.client_max_size):
+	"""
+	Downloads the file at the given url while observing file size and
+	timeout limitations.
+	"""
+	requests_kwargs = {
+		'stream': True,
+		'headers': {'User-Agent': "Steelbea.me LTD needs YOUR files."},
+		'timeout': timeout,
+		'verify': True
+	}
+	temp = b''
+	with requests.get(url, **requests_kwargs) as r:
+		size = 0
+		start_time = time.time()
+		for chunk in r.iter_content(102400):
+			if time.time() - start_time > timeout:
+				raise ValueError('timeout reached')
+			if len(temp) > max_file_size:
+				raise ValueError('response too large')
+			temp += chunk
+
+		if r.headers.get('Content-Disposition'):
+			fname = re.search(r'filename="(.+)"',
+					r.headers['content-disposition'])
+		else:
+			fname = os.path.basename(url)
+	return (fname, temp)
+
+
 async def init_app():
 	"""Initializes the application."""
 	#app = web.Application(middlewares=[buckler_aiohttp.buckler_session])
@@ -108,7 +169,7 @@ async def init_app():
 
 	app.router.add_routes(routes)
 
-	app_wrap = web.Application()
+	app_wrap = web.Application(client_max_size=config.client_max_size)
 	app_wrap.add_subapp(config.url_prefix, app)
 	return app_wrap
 
diff --git a/templates/index.html b/templates/index.html
index fe1e78a..b40c31b 100644
--- a/templates/index.html
+++ b/templates/index.html
@@ -14,6 +14,7 @@
 			<section>
 				<form method="POST" enctype="multipart/form-data">
 					<input type="file" name="files" multiple><br>
+					<input type="url" name="url" placeholder="URL"><br>
 					<input type="checkbox" name="rand_name" id="rand_name" checked>
 					<label for="rand_name">Generate random filename.</label><br>
 					<input type="hidden" name="reponse_type" value="html">