def test_transfer_message_files(self) -> None:
bucket = create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)[0]
hamlet = self.example_user('hamlet')
othello = self.example_user('othello')
upload_message_file('dummy1.txt', len(b'zulip1!'), 'text/plain', b'zulip1!', hamlet)
upload_message_file('dummy2.txt', len(b'zulip2!'), 'text/plain', b'zulip2!', othello)
transfer_message_files_to_s3(1)
attachments = Attachment.objects.all()
self.assertEqual(len(bucket.get_all_keys()), 2)
self.assertEqual(bucket.get_key(attachments[0].path_id).get_contents_as_string(), b'zulip1!')
self.assertEqual(bucket.get_key(attachments[1].path_id).get_contents_as_string(), b'zulip2!')
def test_transfer_message_files(self) -> None:
bucket = create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)[0]
hamlet = self.example_user('hamlet')
othello = self.example_user('othello')
upload_message_file(u'dummy1.txt', len(b'zulip1!'), u'text/plain', b'zulip1!', hamlet)
upload_message_file(u'dummy2.txt', len(b'zulip2!'), u'text/plain', b'zulip2!', othello)
transfer_message_files_to_s3(1)
attachments = Attachment.objects.all()
self.assertEqual(len(bucket.get_all_keys()), 2)
self.assertEqual(bucket.get_key(attachments[0].path_id).get_contents_as_string(), b'zulip1!')
self.assertEqual(bucket.get_key(attachments[1].path_id).get_contents_as_string(), b'zulip2!')
def test_attachment_and_emoji(self) -> None:
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain',
b'zulip!', user_profile)
path_id = url.replace('/user_uploads/', '')
claim_attachment(user_profile=user_profile,
path_id=path_id,
message=message,
is_message_realm_public=True)
realm = Realm.objects.get(string_id='zulip')
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], path_id)
fn = os.path.join(full_data['uploads_dir'], path_id)
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')
fn = os.path.join(
full_data['emoji_dir'],
RealmEmoji.PATH_ID_TEMPLATE.format(realm_id=realm.id,
emoji_file_name='1.png'))
fn = fn.replace('1.png', '')
self.assertEqual('1.png', os.listdir(fn)[0])
def _setup_export_files(self) -> Tuple[str, str, str, bytes]:
realm = Realm.objects.get(string_id='zulip')
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
attachment_path_id = url.replace('/user_uploads/', '')
claim_attachment(
user_profile=user_profile,
path_id=attachment_path_id,
message=message,
is_message_realm_public=True
)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
realm_id=realm.id,
emoji_file_name='1.png',
)
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
message.sender.avatar_source = 'U'
message.sender.save()
return attachment_path_id, emoji_path, original_avatar_path_id, test_image
def extract_and_upload_attachments(message: message.Message, realm: Realm) -> str:
user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT)
attachment_links = []
payload = message.get_payload()
if not isinstance(payload, list):
# This is not a multipart message, so it can't contain attachments.
return ""
for part in payload:
content_type = part.get_content_type()
filename = part.get_filename()
if filename:
attachment = part.get_payload(decode=True)
if isinstance(attachment, bytes):
s3_url = upload_message_file(filename, len(attachment), content_type,
attachment,
user_profile,
target_realm=realm)
formatted_link = "[%s](%s)" % (filename, s3_url)
attachment_links.append(formatted_link)
else:
logger.warning("Payload is not bytes (invalid attachment %s in message from %s)." %
(filename, message.get("From")))
return "\n".join(attachment_links)
def extract_and_upload_attachments(message: message.Message,
realm: Realm) -> str:
user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT)
attachment_links = []
for part in message.walk():
content_type = part.get_content_type()
encoded_filename = part.get_filename()
if not encoded_filename:
continue
filename = handle_header_content(encoded_filename)
if filename:
attachment = part.get_payload(decode=True)
if isinstance(attachment, bytes):
s3_url = upload_message_file(filename,
len(attachment),
content_type,
attachment,
user_profile,
target_realm=realm)
formatted_link = "[%s](%s)" % (filename, s3_url)
attachment_links.append(formatted_link)
else:
logger.warning(
"Payload is not bytes (invalid attachment %s in message from %s)."
% (filename, message.get("From")))
return '\n'.join(attachment_links)
def extract_and_upload_attachments(message: EmailMessage, realm: Realm) -> str:
user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT)
attachment_links = []
for part in message.walk():
content_type = part.get_content_type()
filename = part.get_filename()
if filename:
attachment = part.get_payload(decode=True)
if isinstance(attachment, bytes):
s3_url = upload_message_file(
filename,
len(attachment),
content_type,
attachment,
user_profile,
target_realm=realm,
)
formatted_link = f"[{filename}]({s3_url})"
attachment_links.append(formatted_link)
else:
logger.warning(
"Payload is not bytes (invalid attachment %s in message from %s).",
filename,
message.get("From"),
)
return "\n".join(attachment_links)
def _setup_export_files(self) -> Tuple[str, str, str, bytes]:
realm = Realm.objects.get(string_id='zulip')
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain',
b'zulip!', user_profile)
attachment_path_id = url.replace('/user_uploads/', '')
claim_attachment(user_profile=user_profile,
path_id=attachment_path_id,
message=message,
is_message_realm_public=True)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
realm_id=realm.id,
emoji_file_name='1.png',
)
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
message.sender.avatar_source = 'U'
message.sender.save()
return attachment_path_id, emoji_path, original_avatar_path_id, test_image
def extract_and_upload_attachments(message: message.Message, realm: Realm) -> str:
user_profile = get_system_bot(settings.EMAIL_GATEWAY_BOT)
attachment_links = []
payload = message.get_payload()
if not isinstance(payload, list):
# This is not a multipart message, so it can't contain attachments.
return ""
for part in payload:
content_type = part.get_content_type()
filename = part.get_filename()
if filename:
attachment = part.get_payload(decode=True)
if isinstance(attachment, bytes):
s3_url = upload_message_file(filename, len(attachment), content_type,
attachment,
user_profile,
target_realm=realm)
formatted_link = "[%s](%s)" % (filename, s3_url)
attachment_links.append(formatted_link)
else:
logger.warning("Payload is not bytes (invalid attachment %s in message from %s)." %
(filename, message.get("From")))
return "\n".join(attachment_links)
def test_transfer_message_files(self) -> None:
bucket = create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)[0]
hamlet = self.example_user("hamlet")
othello = self.example_user("othello")
upload_message_file("dummy1.txt", len(b"zulip1!"), "text/plain", b"zulip1!", hamlet)
upload_message_file("dummy2.txt", len(b"zulip2!"), "text/plain", b"zulip2!", othello)
with self.assertLogs(level="INFO"):
transfer_message_files_to_s3(1)
attachments = Attachment.objects.all().order_by("id")
self.assert_length(list(bucket.objects.all()), 2)
self.assertEqual(bucket.Object(attachments[0].path_id).get()["Body"].read(), b"zulip1!")
self.assertEqual(bucket.Object(attachments[1].path_id).get()["Body"].read(), b"zulip2!")
def remove_attachment() -> Dict[str, object]:
user_profile = helpers.example_user("iago")
url = upload_message_file("dummy.txt", len(b"zulip!"), "text/plain",
b"zulip!", user_profile)
attachment_id = url.replace("/user_uploads/", "").split("/")[0]
return {"attachment_id": attachment_id}
def test_get_realm_for_filename(self) -> None:
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
conn.create_bucket(settings.S3_AUTH_UPLOADS_BUCKET)
user_profile = self.example_user('hamlet')
uri = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
path_id = re.sub('/user_uploads/', '', uri)
self.assertEqual(user_profile.realm_id, get_realm_for_filename(path_id))
def upload_attachment(attachment_path: str, sender: UserProfile) -> str:
with open(attachment_path, "rb") as infile:
uploaded_file = File(infile)
uploaded_file_name = os.path.basename(uploaded_file.name)
uploaded_file_size = uploaded_file.size
file_data = uploaded_file.read()
uploaded_uri = upload_message_file(uploaded_file_name,
uploaded_file_size, None, file_data,
sender)
return "\n[" + uploaded_file_name + "](http://localhost:9991" + uploaded_uri + ")\n"
def test_export_files_from_s3(self) -> None:
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
conn.create_bucket(settings.S3_AUTH_UPLOADS_BUCKET)
conn.create_bucket(settings.S3_AVATAR_BUCKET)
realm = Realm.objects.get(string_id='zulip')
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain',
b'zulip!', user_profile)
attachment_path_id = url.replace('/user_uploads/', '')
claim_attachment(user_profile=user_profile,
path_id=attachment_path_id,
message=message,
is_message_realm_public=True)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
realm_id=realm.id,
emoji_file_name='1.png',
)
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], attachment_path_id)
# Test uploads
fields = attachment_path_id.split('/')
fn = os.path.join(full_data['uploads_dir'],
os.path.join(fields[1], fields[2]))
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')
# Test emojis
fn = os.path.join(full_data['emoji_dir'], emoji_path)
fn = fn.replace('1.png', '')
self.assertIn('1.png', os.listdir(fn))
# Test avatars
fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id)
fn_data = open(fn, 'rb').read()
self.assertEqual(fn_data, test_image)
def test_export_files_from_s3(self) -> None:
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
conn.create_bucket(settings.S3_AUTH_UPLOADS_BUCKET)
conn.create_bucket(settings.S3_AVATAR_BUCKET)
realm = Realm.objects.get(string_id='zulip')
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
attachment_path_id = url.replace('/user_uploads/', '')
claim_attachment(
user_profile=user_profile,
path_id=attachment_path_id,
message=message,
is_message_realm_public=True
)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
emoji_path = RealmEmoji.PATH_ID_TEMPLATE.format(
realm_id=realm.id,
emoji_file_name='1.png',
)
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], attachment_path_id)
# Test uploads
fields = attachment_path_id.split('/')
fn = os.path.join(full_data['uploads_dir'], os.path.join(fields[1], fields[2]))
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')
# Test emojis
fn = os.path.join(full_data['emoji_dir'], emoji_path)
fn = fn.replace('1.png', '')
self.assertIn('1.png', os.listdir(fn))
# Test avatars
fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id)
fn_data = open(fn, 'rb').read()
self.assertEqual(fn_data, test_image)
def test_file_upload_local(self) -> None:
user_profile = self.example_user('hamlet')
uri = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
base = '/user_uploads/'
self.assertEqual(base, uri[:len(base)])
path_id = re.sub('/user_uploads/', '', uri)
file_path = os.path.join(settings.LOCAL_UPLOADS_DIR, 'files', path_id)
self.assertTrue(os.path.isfile(file_path))
uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id)
self.assertEqual(len(b'zulip!'), uploaded_file.size)
def test_transfer_message_files(self) -> None:
bucket = create_s3_buckets(settings.S3_AUTH_UPLOADS_BUCKET)[0]
hamlet = self.example_user('hamlet')
othello = self.example_user('othello')
upload_message_file('dummy1.txt', len(b'zulip1!'), 'text/plain',
b'zulip1!', hamlet)
upload_message_file('dummy2.txt', len(b'zulip2!'), 'text/plain',
b'zulip2!', othello)
with self.assertLogs(level="INFO"):
transfer_message_files_to_s3(1)
attachments = Attachment.objects.all().order_by("id")
self.assertEqual(len(list(bucket.objects.all())), 2)
self.assertEqual(
bucket.Object(attachments[0].path_id).get()['Body'].read(),
b'zulip1!')
self.assertEqual(
bucket.Object(attachments[1].path_id).get()['Body'].read(),
b'zulip2!')
def test_export_files_from_local(self) -> None:
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain',
b'zulip!', user_profile)
path_id = url.replace('/user_uploads/', '')
claim_attachment(user_profile=user_profile,
path_id=path_id,
message=message,
is_message_realm_public=True)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
realm = Realm.objects.get(string_id='zulip')
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
message.sender.avatar_source = 'U'
message.sender.save()
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], path_id)
# Test uploads
fn = os.path.join(full_data['uploads_dir'], path_id)
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')
# Test emojis
fn = os.path.join(
full_data['emoji_dir'],
RealmEmoji.PATH_ID_TEMPLATE.format(realm_id=realm.id,
emoji_file_name='1.png'))
fn = fn.replace('1.png', '')
self.assertEqual('1.png', os.listdir(fn)[0])
# Test avatars
fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id)
fn_data = open(fn, 'rb').read()
self.assertEqual(fn_data, test_image)
def test_export_files_from_local(self) -> None:
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
path_id = url.replace('/user_uploads/', '')
claim_attachment(
user_profile=user_profile,
path_id=path_id,
message=message,
is_message_realm_public=True
)
avatar_path_id = user_avatar_path(user_profile)
original_avatar_path_id = avatar_path_id + ".original"
realm = Realm.objects.get(string_id='zulip')
with get_test_image_file('img.png') as img_file:
upload_emoji_image(img_file, '1.png', user_profile)
with get_test_image_file('img.png') as img_file:
upload_avatar_image(img_file, user_profile, user_profile)
test_image = open(get_test_image_file('img.png').name, 'rb').read()
message.sender.avatar_source = 'U'
message.sender.save()
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], path_id)
# Test uploads
fn = os.path.join(full_data['uploads_dir'], path_id)
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')
# Test emojis
fn = os.path.join(full_data['emoji_dir'],
RealmEmoji.PATH_ID_TEMPLATE.format(realm_id=realm.id, emoji_file_name='1.png'))
fn = fn.replace('1.png', '')
self.assertEqual('1.png', os.listdir(fn)[0])
# Test avatars
fn = os.path.join(full_data['avatar_dir'], original_avatar_path_id)
fn_data = open(fn, 'rb').read()
self.assertEqual(fn_data, test_image)
def test_file_upload_s3(self) -> None:
conn = S3Connection(settings.S3_KEY, settings.S3_SECRET_KEY)
bucket = conn.create_bucket(settings.S3_AUTH_UPLOADS_BUCKET)
user_profile = self.example_user('hamlet')
uri = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
base = '/user_uploads/'
self.assertEqual(base, uri[:len(base)])
path_id = re.sub('/user_uploads/', '', uri)
content = bucket.get_key(path_id).get_contents_as_string()
self.assertEqual(b"zulip!", content)
uploaded_file = Attachment.objects.get(owner=user_profile, path_id=path_id)
self.assertEqual(len(b"zulip!"), uploaded_file.size)
self.subscribe(self.example_user("hamlet"), "Denmark")
body = "First message ...[zulip.txt](http://localhost:9991" + uri + ")"
self.send_stream_message(self.example_email("hamlet"), "Denmark", body, "test")
self.assertIn('title="dummy.txt"', self.get_last_message().rendered_content)
def test_attachment(self) -> None:
message = Message.objects.all()[0]
user_profile = message.sender
url = upload_message_file(u'dummy.txt', len(b'zulip!'), u'text/plain', b'zulip!', user_profile)
path_id = url.replace('/user_uploads/', '')
claim_attachment(
user_profile=user_profile,
path_id=path_id,
message=message,
is_message_realm_public=True
)
realm = Realm.objects.get(string_id='zulip')
full_data = self._export_realm(realm)
data = full_data['attachment']
self.assertEqual(len(data['zerver_attachment']), 1)
record = data['zerver_attachment'][0]
self.assertEqual(record['path_id'], path_id)
fn = os.path.join(full_data['uploads_dir'], path_id)
with open(fn) as f:
self.assertEqual(f.read(), 'zulip!')