def test_zip(self):
filename = 'test.zip'
expected = 'test'
actual = archive_basename(filename)
assert expected == actual
def test_tar_gz(self):
filename = 'test.tar.gz'
expected = 'test'
actual = archive_basename(filename)
assert expected == actual
def test_not_allowed_ext(self):
filename = 'test.pdf'
expected = None
actual = archive_basename(filename)
assert expected == actual
def _unique_filename(filename, from_addr='', date='', count=0):
if filename.startswith('='):
charset, encoding, encoded_filename = filename.split('?')[1:4]
if encoding in ('b', 'B'):
filename = base64.b64decode(encoded_filename).decode(charset)
elif encoding in ('q', 'Q'):
filename = encoded_filename
match = QUOTED_PRINTABLE_ENCODED_REGEX.search(filename)
while match:
filename = filename[:match.start()] + chr(
int(match.group()[1:], 16)) + filename[match.end():]
match = QUOTED_PRINTABLE_ENCODED_REGEX.search(filename)
filename = filename.strip().replace(' ', '_')
ext = archive_extension(filename)
basename = archive_basename(filename)
hash_obj = hashlib.sha256()
hash_obj.update(from_addr.encode('utf-8'))
hash_obj.update(date.encode('utf-8'))
hash_obj.update(str(count).encode('utf-8'))
hash_val = hash_obj.hexdigest()[:6]
return '{basename}_{hash_val}{ext}'.format(basename=basename,
hash_val=hash_val,
ext=ext).lower()
def download_attachements(self, directory='.'):
if not os.path.exists(directory):
os.mkdir(directory)
for mail in self.fetch_messages():
count = 0
for part in mail.walk():
count += 1
if (part.get_content_maintype() in ('multipart', 'text')
or part.get('Content-Disposition') is None):
continue
orig_filename = part.get_filename()
if not orig_filename:
print(
term.red('Could not get filename from {} {}'.format(
mail.get('From', ''), mail.get('Subject', ''))))
continue
try:
uniq_filename = self._unique_filename(
orig_filename,
from_addr=mail.get('From', ''),
date=mail.get('Date', ''),
count=count)
except BadExtension as e:
print(term.red(str(e)))
continue
full_path = os.path.join(directory, uniq_filename)
if (not os.path.exists(full_path) and not os.path.exists(
os.path.join(directory,
archive_basename(uniq_filename)))):
print('Writing {}'.format(term.blue(uniq_filename)))
with open(full_path, 'wb') as f:
f.write(part.get_payload(decode=True))
def test_not_allowed_ext(self):
filename = 'test.pdf'
with pytest.raises(BadExtension):
archive_basename(filename)