|
| 1 | + |
| 2 | +import base64 |
| 3 | +import hashlib |
| 4 | +import random |
| 5 | +import re |
| 6 | +import string |
| 7 | + |
| 8 | +from email.mime.image import MIMEImage |
| 9 | + |
| 10 | +def random_string(length, case="lowercase"): |
| 11 | + return "".join(random.choices(getattr(string, f"ascii_{case}") + string.digits, k=length)) |
| 12 | + |
| 13 | +def convert_base64_images(body, attachments): |
| 14 | + |
| 15 | + def repl(match): |
| 16 | + # Capture subtype in case MIMEImage's use of imghdr.what bugs out in guesesing the file type |
| 17 | + subtype = match.group("subtype") |
| 18 | + key = hashlib.md5(base64.b64decode(match.group("data"))).hexdigest().replace("-", "") |
| 19 | + if key not in base64_images: |
| 20 | + base64_images[key] = { |
| 21 | + "data": match.group("data"), |
| 22 | + "subtype": subtype, |
| 23 | + } |
| 24 | + return ' src="cid:image-%s"' % key |
| 25 | + |
| 26 | + # Compile pattern for base64 inline images |
| 27 | + RE_BASE64_SRC = re.compile( |
| 28 | + r' src="data:image/(?P<subtype>gif|png|jpeg|bmp|webp)(?:;charset=utf-8)?;base64,(?P<data>[A-Za-z0-9|+ /]+={0,2})"', |
| 29 | + re.MULTILINE) |
| 30 | + |
| 31 | + base64_images = {} |
| 32 | + |
| 33 | + # Replace and add base64 data to base64_images via repl |
| 34 | + body = re.sub(RE_BASE64_SRC, repl, body) |
| 35 | + for key, image_data in base64_images.items(): |
| 36 | + try: |
| 37 | + image = MIMEImage(base64.b64decode(image_data["data"])) |
| 38 | + except TypeError: |
| 39 | + # Check for subtype if checking fails |
| 40 | + if image_data["subtype"]: |
| 41 | + image = MIMEImage( |
| 42 | + base64.b64decode(image_data["data"]), |
| 43 | + _subtype=image_data["subtype"] |
| 44 | + ) |
| 45 | + else: |
| 46 | + raise |
| 47 | + image.add_header('Content-ID', '<image-%s>' % key) |
| 48 | + image.add_header('Content-Disposition', "inline; filename=%s" % f'image_{random_string(length=6)}') |
| 49 | + attachments.append(image) |
| 50 | + |
| 51 | + return body, attachments |
0 commit comments