Support getting avatars from discord and fedi APIs

.gitignore

@@ -11,5 +11,6 @@ cache/
 __pycache__/
 
 options.json
+secrets.json
 next_serial.txt
 mailer_id.txt

format.py

@@ -2,10 +2,12 @@
 import argparse
 import base64
 import csv
+import hashlib
 import json
 import os
 import string
 import typing
+import urllib.parse
 import xml.etree.ElementTree as ET
 import requests
 import imb
@@ -20,25 +22,73 @@ def iso_code(s: str) -> str:
     return s
 
 
-def get_orig_avatar(url: str, name: str) -> typing.Optional[bytes]:
+def get_discord_avatar(
+    url: urllib.parse.ParseResult, secrets: dict
+) -> (typing.Optional[str], str):
+    try:
+        token = secrets["discord_token"]
+        user_info = requests.get(
+            f"https://discord.com/api/users/{url.path}",
+            headers={"Authorization": f"Bot {token}"},
+        ).json()
+        avatar_hash = user_info["avatar"]
+        return (
+            f"https://cdn.discordapp.com/avatars/{url.path}/{avatar_hash}.png?size=4096",
+            "png",
+        )
+    except KeyError:
+        return None, ""
+
+
+def get_fedi_avatar(
+    url: urllib.parse.ParseResult, secrets: dict
+) -> (typing.Optional[str], str):
+    try:
+        mastodon_api = secrets["mastodon_api"]
+        user_info = requests.get(
+            f"{mastodon_api}/api/v1/accounts/lookup", params={"acct": url.path}
+        ).json()
+        avatar_url = user_info["avatar_static"]
+        extension = avatar_url.split(".")[-1]
+        return avatar_url, extension
+    except KeyError:
+        return None, ""
+
+
+def get_orig_avatar(url: str, basename: str, secrets: dict) -> typing.Optional[bytes]:
     if not os.path.exists("cache"):
         os.mkdir("cache")
-    if os.path.exists("cache/" + name):
-        with open("cache/" + name, "rb") as infile:
+
+    url_parts = urllib.parse.urlparse(url)
+    if url_parts.scheme == "fedi":
+        real_url, extension = get_fedi_avatar(url_parts, secrets)
+    elif url_parts.scheme == "discord":
+        real_url, extension = get_discord_avatar(url_parts, secrets)
+    else:
+        real_url = url
+        extension = url_parts.path.rsplit(".", 1)[1]
+
+    if real_url is None:
+        return None
+
+    img_name = f"cache/{basename}.{extension}"
+
+    if os.path.exists(img_name):
+        with open(img_name, "rb") as infile:
             return infile.read()
-    result = requests.get(url)
+    result = requests.get(real_url)
     if result.ok:
-        with open("cache/" + name, "wb") as outfile:
+        with open(img_name, "wb") as outfile:
             outfile.write(result.content)
         return result.content
     return None
 
 
-def get_avatar(url: str) -> str:
-    name = url.split("?")[0].split("/")[-1]
-    if os.path.exists(f"cache/{name}.svg"):
-        return f"cache/{name}.svg"
-    avatar_raster = get_orig_avatar(url, name)
+def get_avatar(url: str, secrets: dict) -> str:
+    basename = hashlib.sha256(url.encode("utf-8")).hexdigest()
+    if os.path.exists(f"cache/{basename}.svg"):
+        return f"cache/{basename}.svg"
+    avatar_raster = get_orig_avatar(url, basename, secrets)
     if avatar_raster is None:
         return ""
 
@@ -50,9 +100,9 @@ def get_avatar(url: str) -> str:
             xlink:href="data:;base64,{base64.b64encode(avatar_raster).decode("utf-8")}" />
     </svg>"""
 
-    with open(f"cache/{name}.svg", "w") as svgfile:
+    with open(f"cache/{basename}.svg", "w") as svgfile:
         svgfile.write(svg_text)
-    return f"cache/{name}.svg"
+    return f"cache/{basename}.svg"
 
 
 def get_country_name(
@@ -111,7 +161,7 @@ parser.add_argument(
     "-n",
     "--no-content",
     action="store_true",
-    help="Skip content, e.g. to make postcard back labels"
+    help="Skip content, e.g. to make postcard back labels",
 )
 
 args = parser.parse_args()
@@ -123,8 +173,11 @@ root = ET.parse(
 csvfile = open(args.address_file)
 rows = csv.DictReader(csvfile)
 
+with open("secrets.json") as secrets_file:
+    secrets = json.load(secrets_file)
+
 current_serial = imb.get_first_serial()
-mid = int(open("mailer_id.txt").read().strip())
+mid = secrets.get("mailer_id")
 
 
 cards = []
@@ -141,7 +194,7 @@ for row in rows:
     address = row["Address"].split("\n") + country
 
     if row.get("Avatar", "") != "":
-        avatar = get_avatar(row["Avatar"])
+        avatar = get_avatar(row["Avatar"], secrets)
     else:
         avatar = None
 
@@ -160,7 +213,7 @@ serial = imb.get_first_serial()
 if args.origin == "us":
     for card in cards:
         dpc = card["row"].get("DPC", "")
-        if dpc != "":
+        if dpc != "" and mid is not None:
             card["imb"] = imb.generate(
                 0, 310, mid, serial, dpc.replace(" ", "").replace("-", "")
             )