154 lines
6.7 KiB
Python
154 lines
6.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate a cover image for a THREATMARKET intel pack via ComfyUI SDXL.
|
|
|
|
Zero-failure design: tries ComfyUI (juggernautXL, then dreamshaperXL, then
|
|
sdxl-base), then falls back to a locally-rendered SVG->PNG-style placeholder
|
|
(a real, deterministic dark radar-brand cover - not fake data, just brand art).
|
|
Always writes /opt/threatmarket/static/covers/<pack_id>.png (or .svg) and
|
|
returns the web path, or None only if the filesystem itself is unwritable.
|
|
|
|
Usage: from cover_gen import generate_cover
|
|
path = generate_cover("chewy.com", "2026-09-25")
|
|
"""
|
|
import base64
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
|
|
COMFY = os.environ.get("COMFY_URL", "http://10.30.20.29:8188")
|
|
COVER_DIR = os.environ.get("COVER_DIR", "/opt/threatmarket/static/covers")
|
|
WEB_PREFIX = "/covers"
|
|
# Order = preference. Lightning models = few steps = fast + reliable.
|
|
CKPTS = ["flux1-schnell-fp8.safetensors",
|
|
"juggernautXL_ragnarok.safetensors",
|
|
"dreamshaperXL_lightningDPMSDE.safetensors"]
|
|
TIMEOUT_SUBMIT = 20
|
|
TIMEOUT_POLL = 150
|
|
TIMEOUT_FETCH = 30
|
|
|
|
|
|
def _prompt_text(domain, date):
|
|
return ('dark hacker operations room at night, glowing green matrix rain code '
|
|
'walls, multiple monitors showing world-map network dashboards and '
|
|
'breach data, empty desk chair, dramatic cyan and green lighting, '
|
|
'cinematic depth of field, gritty cyberpunk atmosphere, wide dark '
|
|
'area at top for title text, professional threat-intelligence dossier '
|
|
'cover art, no text, no letters, no words, 4k')
|
|
|
|
|
|
def _workflow(ckpt, seed):
|
|
return {
|
|
"3": {"class_type": "KSampler", "inputs": {
|
|
"seed": seed, "steps": 20, "cfg": 1.8, "sampler_name": "euler",
|
|
"scheduler": "sgm_uniform", "denoise": 1.0,
|
|
"model": ["4", 0], "positive": ["6", 0], "negative": ["7", 0],
|
|
"latent_image": ["5", 0]}},
|
|
"4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": ckpt}},
|
|
"5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 640, "batch_size": 1}},
|
|
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "placeholder", "clip": ["4", 1]}},
|
|
"7": {"class_type": "CLIPTextEncode", "inputs": {"text": "text, watermark, blurry, low quality, oversaturated", "clip": ["4", 1]}},
|
|
"8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}},
|
|
"9": {"class_type": "SaveImage", "inputs": {"filename_prefix": "tmcover", "images": ["8", 0]}},
|
|
}
|
|
|
|
|
|
def _comfy_generate(domain, date, seed):
|
|
"""Try each checkpoint; return PNG bytes or None."""
|
|
for ckpt in CKPTS:
|
|
try:
|
|
wf = _workflow(ckpt, seed)
|
|
# fix node 6 text properly (dict-literal scoping hack above)
|
|
wf["6"]["inputs"]["text"] = _prompt_text(domain, date)
|
|
req = urllib.request.Request(
|
|
COMFY + "/prompt",
|
|
data=json.dumps({"prompt": wf}).encode(),
|
|
headers={"Content-Type": "application/json"})
|
|
with urllib.request.urlopen(req, timeout=TIMEOUT_SUBMIT) as r:
|
|
resp = json.loads(r.read())
|
|
pid = resp.get("prompt_id")
|
|
if not pid:
|
|
continue
|
|
import time
|
|
deadline = time.time() + TIMEOUT_POLL
|
|
while time.time() < deadline:
|
|
with urllib.request.urlopen(COMFY + "/history/" + pid, timeout=15) as r:
|
|
hist = json.loads(r.read())
|
|
if pid in hist and hist[pid].get("outputs"):
|
|
outs = hist[pid]["outputs"]
|
|
for node in outs.values():
|
|
for img in node.get("images", []):
|
|
fn = img["filename"]
|
|
sub = img.get("subfolder", "")
|
|
imgtype = img.get("type", "output")
|
|
url = f"{COMFY}/view?filename={fn}&subfolder={sub}&type={imgtype}"
|
|
with urllib.request.urlopen(url, timeout=TIMEOUT_FETCH) as fr:
|
|
return fr.read()
|
|
time.sleep(3)
|
|
except Exception:
|
|
continue
|
|
return None
|
|
|
|
|
|
def _fallback_svg(domain, date):
|
|
"""Deterministic branded SVG cover - always succeeds if fs writable."""
|
|
seed = sum(ord(c) for c in domain) % 360
|
|
svg = f"""<svg xmlns='http://www.w3.org/2000/svg' width='1024' height='640'>
|
|
<defs><radialGradient id='g' cx='75%' cy='20%'><stop offset='0%' stop-color='hsl({seed},80%,22%)'/>
|
|
<stop offset='100%' stop-color='#0a0e1a'/></radialGradient></defs>
|
|
<rect width='1024' height='640' fill='url(#g)'/>
|
|
<g stroke='hsl({seed},90%,60%)' stroke-opacity='0.25' fill='none'>"""
|
|
for i in range(6):
|
|
svg += f"<circle cx='780' cy='200' r='{40+i*45}' stroke-width='1'/>"
|
|
svg += f"</g><g fill='hsl({seed},90%,65%)'>"
|
|
import math
|
|
for i in range(14):
|
|
x = 120 + (i * 61) % 700
|
|
y = 420 + (i * 37) % 160
|
|
r = 2 + (i % 3)
|
|
svg += f"<circle cx='{x}' cy='{y}' r='{r}' opacity='{0.35 + (i % 5) * 0.12}'/>"
|
|
svg += f"""</g>
|
|
<text x='48' y='90' font-family='monospace' font-size='30' fill='#e8f4ff' font-weight='bold'>INTEL PACK</text>
|
|
<text x='48' y='130' font-family='monospace' font-size='22' fill='hsl({seed},90%,70%)'>{domain}</text>
|
|
<text x='48' y='162' font-family='monospace' font-size='16' fill='#8fa3bf'>{date} — THREATMARKET</text>
|
|
</svg>"""
|
|
return svg.encode()
|
|
|
|
|
|
def generate_cover(domain, date):
|
|
"""Returns web path of cover image, or None. Never raises."""
|
|
os.makedirs(COVER_DIR, exist_ok=True)
|
|
pid = (domain + "-" + date.replace("-", "")).replace("/", "_")
|
|
png_path = os.path.join(COVER_DIR, pid + ".png")
|
|
svg_path = os.path.join(COVER_DIR, pid + ".svg")
|
|
if os.path.exists(png_path):
|
|
return f"{WEB_PREFIX}/{pid}.png"
|
|
if os.path.exists(svg_path):
|
|
return f"{WEB_PREFIX}/{pid}.svg"
|
|
seed = sum(ord(c) for c in pid) % (2 ** 31)
|
|
png = _comfy_generate(domain, date, seed)
|
|
if png:
|
|
try:
|
|
with open(png_path, "wb") as f:
|
|
f.write(png)
|
|
try:
|
|
from text_stamp import stamp
|
|
stamp(png_path, domain)
|
|
except Exception:
|
|
pass # art still fine even if stamp fails
|
|
return f"{WEB_PREFIX}/{pid}.png"
|
|
except OSError:
|
|
pass
|
|
try:
|
|
with open(svg_path, "wb") as f:
|
|
f.write(_fallback_svg(domain, date))
|
|
return f"{WEB_PREFIX}/{pid}.svg"
|
|
except OSError:
|
|
return None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
d = sys.argv[1] if len(sys.argv) > 1 else "mozilla.org"
|
|
dt = sys.argv[2] if len(sys.argv) > 2 else "2026-09-24"
|
|
print(generate_cover(d, dt) or "FAILED")
|