# HG changeset patch # User Brian Neal # Date 1444093664 18000 # Node ID f5aa74dcdd7a0797ed7b155ff4b0d4b076e15096 # Parent 8c3d52b7cbd1ca98b2704767ecd74bb3cb591934 Ensure temporary files get deleted during hotlinking. diff -r 8c3d52b7cbd1 -r f5aa74dcdd7a core/download.py --- a/core/download.py Mon Oct 05 20:06:24 2015 -0500 +++ b/core/download.py Mon Oct 05 20:07:44 2015 -0500 @@ -49,6 +49,7 @@ shutil.copyfileobj(r.raw, fp) except requests.RequestException: logger.exception("download_file download exception") + os.remove(path) raise file_size = os.stat(path).st_size diff -r 8c3d52b7cbd1 -r f5aa74dcdd7a core/functions.py --- a/core/functions.py Mon Oct 05 20:06:24 2015 -0500 +++ b/core/functions.py Mon Oct 05 20:07:44 2015 -0500 @@ -1,4 +1,5 @@ """This file houses various core utility functions""" +from contextlib import contextmanager import datetime import logging import os @@ -31,6 +32,15 @@ os.remove(self.filename) +@contextmanager +def remove_file(path): + """Context manager for removing a file when the context is exited.""" + try: + yield path + finally: + os.remove(path) + + def send_mail(subject, message, from_email, recipient_list, reply_to=None, defer=True): """ The main send email function. Use this function to send email from the diff -r 8c3d52b7cbd1 -r f5aa74dcdd7a user_photos/forms.py --- a/user_photos/forms.py Mon Oct 05 20:06:24 2015 -0500 +++ b/user_photos/forms.py Mon Oct 05 20:07:44 2015 -0500 @@ -8,7 +8,7 @@ from django.conf import settings from core.download import download_file -from core.functions import TemporaryFile +from core.functions import remove_file, TemporaryFile from core.images.upload import upload from core.s3 import S3Bucket from core.services import get_redis_connection @@ -150,20 +150,20 @@ return url # Try to download the file - path = download_file(url) + with remove_file(download_file(url)) as path: - # Upload it to our S3 bucket - bucket = S3Bucket(access_key=settings.USER_PHOTOS_ACCESS_KEY, - secret_key=settings.USER_PHOTOS_SECRET_KEY, - base_url=settings.HOT_LINK_PHOTOS_BASE_URL, - bucket_name=settings.HOT_LINK_PHOTOS_BUCKET) + # Upload it to our S3 bucket + bucket = S3Bucket(access_key=settings.USER_PHOTOS_ACCESS_KEY, + secret_key=settings.USER_PHOTOS_SECRET_KEY, + base_url=settings.HOT_LINK_PHOTOS_BASE_URL, + bucket_name=settings.HOT_LINK_PHOTOS_BUCKET) - now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') - metadata = {'user': self.user.username, 'date': now} + now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') + metadata = {'user': self.user.username, 'date': now} - url, _ = upload(filename=path, - bucket=bucket, - metadata=metadata, - new_size=settings.USER_PHOTOS_MAX_SIZE) + url, _ = upload(filename=path, + bucket=bucket, + metadata=metadata, + new_size=settings.USER_PHOTOS_MAX_SIZE) return url diff -r 8c3d52b7cbd1 -r f5aa74dcdd7a user_photos/tests/test_forms.py --- a/user_photos/tests/test_forms.py Mon Oct 05 20:06:24 2015 -0500 +++ b/user_photos/tests/test_forms.py Mon Oct 05 20:07:44 2015 -0500 @@ -1,6 +1,7 @@ """ Tests for the forms in the user_photos application. """ +from contextlib import contextmanager import mock from django import forms @@ -11,6 +12,11 @@ from user_photos.forms import HotLinkImageForm +@contextmanager +def fake_remove_file(path): + yield path + + class HotLinkImageFormTestCase(TestCase): def setUp(self): @@ -41,6 +47,7 @@ @mock.patch('user_photos.forms.rate_limit_user') @mock.patch('user_photos.forms.download_file') + @mock.patch('user_photos.forms.remove_file', new=fake_remove_file) @mock.patch('user_photos.forms.S3Bucket') @mock.patch('user_photos.forms.upload') def test_white_listed_url(self, upload_mock, bucket_mock, dl_mock, rate_limit_mock): @@ -56,6 +63,7 @@ @mock.patch('user_photos.forms.rate_limit_user') @mock.patch('user_photos.forms.download_file') + @mock.patch('user_photos.forms.remove_file', new=fake_remove_file) @mock.patch('user_photos.forms.S3Bucket') @mock.patch('user_photos.forms.upload') def test_happy_path(self, upload_mock, bucket_mock, dl_mock, rate_limit_mock):