zulip/zerver/tests/test_upload_local.py

import os
import re
from io import BytesIO, StringIO
from urllib.parse import urlsplit

from django.conf import settings
from PIL import Image

import zerver.lib.upload
from zerver.lib.avatar_hash import user_avatar_path
from zerver.lib.test_classes import UploadSerializeMixin, ZulipTestCase
from zerver.lib.test_helpers import get_test_image_file, read_test_image_file
from zerver.lib.upload import (
    all_message_attachments,
    delete_export_tarball,
    delete_message_attachment,
    delete_message_attachments,
    save_attachment_contents,
    upload_emoji_image,
    upload_export_tarball,
    upload_message_attachment,
)
from zerver.lib.upload.base import DEFAULT_EMOJI_SIZE, MEDIUM_AVATAR_SIZE, resize_avatar
from zerver.lib.upload.local import write_local_file
from zerver.models import Attachment, RealmEmoji
from zerver.models.realms import get_realm
from zerver.models.users import get_system_bot


class LocalStorageTest(UploadSerializeMixin, ZulipTestCase):
    def test_upload_message_attachment(self) -> None:
        user_profile = self.example_user("hamlet")
        url = upload_message_attachment(
            "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile
        )

        base = "/user_uploads/"
        self.assertEqual(base, url[: len(base)])
        path_id = re.sub("/user_uploads/", "", url)
        assert settings.LOCAL_UPLOADS_DIR is not None
        assert settings.LOCAL_FILES_DIR is not None
        file_path = os.path.join(settings.LOCAL_FILES_DIR, path_id)
        self.assertTrue(os.path.isfile(file_path))

        uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id)
        self.assert_length(b"zulip!", uploaded_file.size)

    def test_save_attachment_contents(self) -> None:
        user_profile = self.example_user("hamlet")
        url = upload_message_attachment(
            "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile
        )

        path_id = re.sub("/user_uploads/", "", url)
        output = BytesIO()
        save_attachment_contents(path_id, output)
        self.assertEqual(output.getvalue(), b"zulip!")

    def test_upload_message_attachment_local_cross_realm_path(self) -> None:
        """
        Verifies that the path of a file uploaded by a cross-realm bot to another
        realm is correct.
        """

        internal_realm = get_realm(settings.SYSTEM_BOT_REALM)
        zulip_realm = get_realm("zulip")
        user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT, internal_realm.id)
        self.assertEqual(user_profile.realm, internal_realm)

        url = upload_message_attachment(
            "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile, zulip_realm
        )
        # Ensure the correct realm id of the target realm is used instead of the bot's realm.
        self.assertTrue(url.startswith(f"/user_uploads/{zulip_realm.id}/"))

    def test_delete_message_attachment(self) -> None:
        self.login("hamlet")
        fp = StringIO("zulip!")
        fp.name = "zulip.txt"
        result = self.client_post("/json/user_uploads", {"file": fp})

        response_dict = self.assert_json_success(result)
        path_id = re.sub("/user_uploads/", "", response_dict["uri"])

        assert settings.LOCAL_FILES_DIR is not None
        file_path = os.path.join(settings.LOCAL_FILES_DIR, path_id)
        self.assertTrue(os.path.isfile(file_path))

        self.assertTrue(delete_message_attachment(path_id))
        self.assertFalse(os.path.isfile(file_path))

    def test_delete_message_attachments(self) -> None:
        assert settings.LOCAL_UPLOADS_DIR is not None
        assert settings.LOCAL_FILES_DIR is not None

        user_profile = self.example_user("hamlet")
        path_ids = []
        for n in range(1, 1005):
            url = upload_message_attachment(
                "dummy.txt", len(b"zulip!"), "text/plain", b"zulip!", user_profile
            )
            base = "/user_uploads/"
            self.assertEqual(base, url[: len(base)])
            path_id = re.sub("/user_uploads/", "", url)
            path_ids.append(path_id)
            file_path = os.path.join(settings.LOCAL_FILES_DIR, path_id)
            self.assertTrue(os.path.isfile(file_path))

        delete_message_attachments(path_ids)
        for path_id in path_ids:
            file_path = os.path.join(settings.LOCAL_FILES_DIR, path_id)
            self.assertFalse(os.path.isfile(file_path))

    def test_all_message_attachments(self) -> None:
        write_local_file("files", "foo", b"content")
        write_local_file("files", "bar/baz", b"content")
        write_local_file("files", "bar/troz", b"content")
        write_local_file("files", "test/other/file", b"content")
        found_files = [r[0] for r in all_message_attachments()]
        self.assertEqual(sorted(found_files), ["bar/baz", "bar/troz", "foo", "test/other/file"])

    def test_avatar_url(self) -> None:
        self.login("hamlet")
        with get_test_image_file("img.png") as image_file:
            result = self.client_post("/json/users/me/avatar", {"file": image_file})

        response_dict = self.assert_json_success(result)
        self.assertIn("avatar_url", response_dict)
        base = "/user_avatars/"
        url = self.assert_json_success(result)["avatar_url"]
        self.assertEqual(base, url[: len(base)])

        # That URL is accessible when logged out
        self.logout()
        result = self.client_get(url)
        self.assertEqual(result.status_code, 200)

        # We get a resized avatar from it
        image_data = read_test_image_file("img.png")
        resized_avatar = resize_avatar(image_data)
        self.assertEqual(resized_avatar, result.getvalue())

        with self.settings(DEVELOPMENT=False):
            # In production, this is an X-Accel-Redirect to the
            # on-disk content, which nginx serves
            result = self.client_get(url)
            self.assertEqual(result.status_code, 200)
            internal_redirect_path = urlsplit(url).path.replace(
                "/user_avatars/", "/internal/local/user_avatars/"
            )
            self.assertEqual(result["X-Accel-Redirect"], internal_redirect_path)
            self.assertEqual(b"", result.content)

    def test_ensure_avatar_image(self) -> None:
        user_profile = self.example_user("hamlet")
        file_path = user_avatar_path(user_profile)

        write_local_file("avatars", file_path + ".original", read_test_image_file("img.png"))

        assert settings.LOCAL_UPLOADS_DIR is not None
        assert settings.LOCAL_AVATARS_DIR is not None
        image_path = os.path.join(settings.LOCAL_AVATARS_DIR, file_path + ".original")
        with open(image_path, "rb") as f:
            image_data = f.read()

        resized_avatar = resize_avatar(image_data)
        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile)
        output_path = os.path.join(settings.LOCAL_AVATARS_DIR, file_path + ".png")
        with open(output_path, "rb") as original_file:
            self.assertEqual(resized_avatar, original_file.read())

        resized_avatar = resize_avatar(image_data, MEDIUM_AVATAR_SIZE)
        zerver.lib.upload.upload_backend.ensure_avatar_image(user_profile, is_medium=True)
        output_path = os.path.join(settings.LOCAL_AVATARS_DIR, file_path + "-medium.png")
        with open(output_path, "rb") as original_file:
            self.assertEqual(resized_avatar, original_file.read())

    def test_get_emoji_url(self) -> None:
        user_profile = self.example_user("hamlet")
        file_name = "emoji.png"

        with get_test_image_file("img.png") as image_file:
            upload_emoji_image(image_file, file_name, user_profile)
        url = zerver.lib.upload.upload_backend.get_emoji_url(file_name, user_profile.realm_id)

        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
            realm_id=user_profile.realm_id,
            emoji_file_name=file_name,
        )
        expected_url = f"/user_avatars/{emoji_path}"
        self.assertEqual(expected_url, url)

        file_name = "emoji.gif"
        with get_test_image_file("animated_img.gif") as image_file:
            upload_emoji_image(image_file, file_name, user_profile)
        url = zerver.lib.upload.upload_backend.get_emoji_url(file_name, user_profile.realm_id)
        still_url = zerver.lib.upload.upload_backend.get_emoji_url(
            file_name, user_profile.realm_id, still=True
        )

        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
            realm_id=user_profile.realm_id,
            emoji_file_name=file_name,
        )

        still_emoji_path = RealmEmoji.STILL_PATH_ID_TEMPLATE.format(
            realm_id=user_profile.realm_id,
            emoji_filename_without_extension=os.path.splitext(file_name)[0],
        )
        expected_url = f"/user_avatars/{emoji_path}"
        self.assertEqual(expected_url, url)
        expected_still_url = f"/user_avatars/{still_emoji_path}"
        self.assertEqual(expected_still_url, still_url)

    def test_emoji_upload(self) -> None:
        user_profile = self.example_user("hamlet")
        file_name = "emoji.png"

        with get_test_image_file("img.png") as image_file:
            upload_emoji_image(image_file, file_name, user_profile)

        emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
            realm_id=user_profile.realm_id,
            emoji_file_name=file_name,
        )

        assert settings.LOCAL_AVATARS_DIR is not None
        file_path = os.path.join(settings.LOCAL_AVATARS_DIR, emoji_path)
        with open(file_path + ".original", "rb") as original_file:
            self.assertEqual(read_test_image_file("img.png"), original_file.read())

        expected_size = (DEFAULT_EMOJI_SIZE, DEFAULT_EMOJI_SIZE)
        with Image.open(file_path) as resized_image:
            self.assertEqual(expected_size, resized_image.size)

    def test_tarball_upload_and_deletion(self) -> None:
        user_profile = self.example_user("iago")
        self.assertTrue(user_profile.is_realm_admin)

        assert settings.TEST_WORKER_DIR is not None
        tarball_path = os.path.join(settings.TEST_WORKER_DIR, "tarball.tar.gz")
        with open(tarball_path, "w") as f:
            f.write("dummy")

        assert settings.LOCAL_AVATARS_DIR is not None
        url = upload_export_tarball(user_profile.realm, tarball_path)
        self.assertTrue(os.path.isfile(os.path.join(settings.LOCAL_AVATARS_DIR, tarball_path)))

        result = re.search(re.compile(r"([A-Za-z0-9\-_]{24})"), url)
        if result is not None:
            random_name = result.group(1)
        expected_url = f"http://zulip.testserver/user_avatars/exports/{user_profile.realm_id}/{random_name}/tarball.tar.gz"
        self.assertEqual(expected_url, url)

        # Delete the tarball.
        with self.assertLogs(level="WARNING") as warn_log:
            self.assertIsNone(delete_export_tarball("/not_a_file"))
        self.assertEqual(
            warn_log.output,
            ["WARNING:root:not_a_file does not exist. Its entry in the database will be removed."],
        )
        path_id = urlsplit(url).path
        self.assertEqual(delete_export_tarball(path_id), path_id)