chore: Delete unused scripts

This commit is contained in:
abdus committed 2026-10-20 07:43:00 +03:00
1 parent 5ecff074af
commit f449d692a8
34 files changed
-6047

No files matched your search

View File
Whitespace-only changes.
View File
Whitespace-only changes.
-48
View File
@@ -1,48 +0,0 @@
import contextlib
from playwright.sync_api import sync_playwright, Page
from time import sleep
@contextlib.contextmanager
def launch_browser() -> Page:
with sync_playwright() as playwright:
with playwright.chromium.launch(headless=False) as browser:
with browser.new_context() as context:
with context.new_page() as page:
yield page
def book_slot():
with launch_browser() as page:
page.goto("https://service.berlin.de/dienstleistung/120686/")
page.get_by_role("complementary").get_by_role("link", name="Termin berlinweit suchen").click()
with page.expect_navigation(url='https://service.berlin.de/terminvereinbarung/termin/time/*', timeout=0):
pass
with page.expect_navigation(url='https://service.berlin.de/terminvereinbarung/termin/register/*', timeout=0):
pass
page.get_by_label("Ihr Vor- und Nachname *").fill("Paulina Gutkowska")
page.get_by_label("Ihre E-Mail Adresse *").fill("gutpaula11@gmail.com")
if page.is_visible("Ihre Telefonnummer oder Mobilfunknummer *"):
page.get_by_label("Ihre Telefonnummer oder Mobilfunknummer *").fill("+48501099723")
page.get_by_role("combobox").select_option("0")
page.get_by_label(
"Ich erkläre mich mit den Nutzungsbedingungen einverstanden und akzeptiere diese. (Pflichtangabe. Termineintrag nur bei Einverständnis möglich) *"
).check()
with page.expect_navigation():
page.get_by_role("button", name="Termin eintragen").click()
sleep(10)
input()
def main():
book_slot()
if __name__ == '__main__':
main()
-75
View File
@@ -1,75 +0,0 @@
import contextlib
import datetime
import locale
import httpx
from bs4 import BeautifulSoup
import berlin
http = httpx.Client(
headers={
'user-agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36'
},
follow_redirects=True,
)
@contextlib.contextmanager
def override_locale(category: int, val: str) -> None:
prev = locale.getlocale(category)
locale.setlocale(category, val)
yield
locale.setlocale(category, prev)
def get_available_slots():
res = http.get('https://service.berlin.de/dienstleistung/120686/')
res.raise_for_status()
soup = BeautifulSoup(res.text, 'html.parser')
link = soup.select_one('[role="complementary"] .zmstermin-multi a')
calendar_url = link.attrs['href']
try:
# redirect fails, we just need the cookies
_ = http.get(calendar_url)
except:
pass
res = http.get('https://service.berlin.de/terminvereinbarung/termin/day/')
res.raise_for_status()
return extract_days(res.text)
def extract_days(html: str) -> list[datetime.date]:
soup = BeautifulSoup(html, 'html.parser')
found = []
for it in soup.select('.calendar-month-table .buchbar'):
month_year = it.find_parent(attrs={'class': 'calendar-month-table'}).select_one('thead .month').text.strip()
day = it.text.strip()
found.append(f'{day} {month_year}')
with override_locale(locale.LC_TIME, 'de_DE'):
days = [datetime.datetime.strptime(it, '%d %B %Y').date() for it in found]
return days
def main():
days = get_available_slots()
good_days = [it for it in days if it < datetime.date(2023, 1, 27)]
if not good_days:
print('no slots')
return
print('found slots')
berlin.book_slot()
for it in days:
print(it)
if __name__ == '__main__':
main()
-156
View File
@@ -1,156 +0,0 @@
import argparse
import asyncio
import contextlib
import dataclasses
import datetime
import logging
import random
import re
import subprocess
import urllib.parse
from pathlib import Path
from playwright.async_api import async_playwright, Browser, BrowserContext, Page
@contextlib.asynccontextmanager
async def launch_browser(headless: bool = False) -> Page:
async with async_playwright() as playwright:
browser = await playwright.chromium.launch(headless=headless)
browser: Browser
async with browser:
ctx = await browser.new_context()
ctx: BrowserContext
async with ctx:
page = await ctx.new_page()
async with page:
yield page
class Error(Exception):
pass
def alert(text: str):
subprocess.run(['open', f'alfred://runtrigger/dev.abdus.integrations/remind/?argument={urllib.parse.quote(text)}'])
dump_dir = Path('~/Desktop').expanduser() / 'immigration'
dump_dir.mkdir(parents=True, exist_ok=True)
@dataclasses.dataclass
class ResidencyApplication:
citizenship_country: str
num_applicants: int
is_living_with_family: bool
residency_category: str
residency_purpose: str
async def check_slots(page: Page, values: ResidencyApplication, retry_count: int = 1) -> list[datetime.datetime]:
await page.add_init_script('''Object.defineProperty(navigator, 'webdriver', { get: () => false })''')
window_id = random.randint(1000, 9999)
req_id = random.randint(0, 999)
await page.goto(
f'https://otv.verwalt-berlin.de/ams/TerminBuchen/wizardng?dswid={window_id}&dsrid={req_id}',
wait_until='networkidle',
)
async with page.expect_navigation():
await page.click('.langBlock [name="txtEn"] a')
await page.check('[name="gelesen"]')
async with page.expect_navigation(url=re.compile('st=2'), wait_until='networkidle'):
await page.click('[name="applicationForm:managedForm:proceed"]')
await page.get_by_role("combobox", name="Citizenship *").select_option(label=values.citizenship_country)
await asyncio.sleep(0.3)
# how many applicants
await page.select_option('[name="personenAnzahl_normal"]', value=str(values.num_applicants))
await asyncio.sleep(0.3)
# living with family
await page.select_option('[name="lebnBrMitFmly"]', value='1' if values.is_living_with_family else '2')
await asyncio.sleep(0.3)
await page.get_by_text("Apply for a residence title").click()
await page.locator("label").filter(has_text=values.residency_category).click()
await page.get_by_text(values.residency_purpose).click()
stop_at = datetime.datetime.now() + datetime.timedelta(minutes=28)
while 'st=2' in page.url:
retry_count -= 1
async with page.expect_navigation(url=re.compile(r'st='), wait_until='networkidle', timeout=60_000):
await page.get_by_role("button", name="Next").click()
if datetime.datetime.now() > stop_at:
raise TimeoutError('could not find a slot')
if not retry_count:
return []
await asyncio.sleep(random.randint(4, 10))
continue
if 'st=3' in page.url:
dates = []
for el in await page.query_selector_all('.ui-datepicker-calendar td:not(.ui-datepicker-unselectable)'):
dates.append(
datetime.date(
year=int(await el.get_attribute('data-year')),
month=int(await el.get_attribute('data-month')),
day=int(await el.inner_text()),
)
)
alert("\n".join(['empty slots found', *[it.isoformat() for it in dates]]))
now = int(datetime.datetime.now().timestamp())
html_path = dump_dir / f'immigration_{now}.html'
html_path.write_text(await page.content())
await page.screenshot(path=dump_dir / f'immigration_{now}.png')
input()
input()
def parse_args():
arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
arger.add_argument(
'--headless', '--silent', '--quiet', action='store_true', dest='headless', help='Run in headless mode'
)
arger.add_argument('--retry-count', '--retry', type=int, default=1, help='Run in headless mode')
args, _ = arger.parse_known_args()
return args
async def main():
args = parse_args()
async with launch_browser(headless=args.headless) as page:
try:
values = ResidencyApplication(
citizenship_country='India',
num_applicants=1,
is_living_with_family=False,
residency_category='Educational purposes',
residency_purpose='Residence permit for the purpose of studying (sect. 16b)',
)
slots = await check_slots(page, values=values, retry_count=args.retry_count)
if not slots:
print('no slots yet')
return
except Exception as e:
logging.exception('got an error')
input()
if __name__ == '__main__':
asyncio.run(main())
logging.basicConfig(level=logging.INFO)
-167
View File
@@ -1,167 +0,0 @@
import csv
import html
import sys
# language=css
page_css = """
body {
margin: 0;
padding: 0.5rem;
font-family: menlo, consolas, monospace;
font-size: 14px;
}
input {
font-size: inherit;
font-family: inherit;
background-color: transparent;
padding: 0.25em 0.5em;
}
table {
border-collapse: collapse;
width: 100%;
}
th, td {
border: 1px solid black;
padding: 8px;
text-align: left;
}
th {
background-color: #f2f2f2;
}
tr:nth-child(even) {
background-color: #fdfdfd;
}
"""
# language=js
page_script = """
function sortableTable(table) {
const sortStates = new Map();
table.querySelectorAll('th').forEach((header, colIndex) => {
header.style.cursor = 'pointer';
header.onclick = () => {
const isAsc = !sortStates.get(header);
const tbody = table.querySelector('tbody');
const rows = [...tbody.rows];
rows.sort((a, b) => {
const [valA, valB] = [a, b].map(row =>
row.cells[colIndex].textContent.trim());
return /^\d+$/.test(valA)
? (valA - valB) * (isAsc ? 1 : -1)
: valA.localeCompare(valB) * (isAsc ? 1 : -1);
});
sortStates.set(header, isAsc);
tbody.replaceChildren(...rows); // Better than innerHTML = ''
};
});
}
function debounce(func, delay) {
let timeoutId;
return function(...args) {
clearTimeout(timeoutId);
timeoutId = setTimeout(() => {
func.apply(this, args);
}, delay);
};
}
function addFilterInput(tableElement) {
const filterInput = document.createElement('input');
filterInput.type = 'text';
filterInput.placeholder = "Search...";
tableElement.parentNode.insertBefore(filterInput, tableElement);
const debouncedFilter = debounce(filter, 100);
filterInput.addEventListener('input', () => {
debouncedFilter();
});
function filter() {
const filterValue = filterInput.value.trim();
const rows = tableElement.querySelectorAll('tbody tr');
rows.forEach(row => {
const rowData = row.textContent.toLowerCase();
if (filterValue.startsWith('/') && filterValue.endsWith('/')) {
const regex = new RegExp(filterValue.slice(1, -1), 'ig');
row.style.display = regex.test(rowData) ? '' : 'none';
} else {
row.style.display = rowData.includes(filterValue.toLowerCase()) ? '' : 'none';
}
});
}
}
sortableTable(document.querySelector('table'));
addFilterInput(document.querySelector('table'));
"""
def csv_to_html(csv_input: str) -> str:
try:
lines = csv_input.splitlines()
delimiter = "\t" if "\t" in lines[0] else ","
reader = csv.reader(lines, delimiter=delimiter)
header = next(reader) # Get the header row
header_html = f'{"".join(f"<th>{html.escape(col)}</th>" for col in header)}'
rows = [
f"<tr>{''.join(f'<td>{html.escape(cell)}</td>' for cell in row)}</tr>"
for row in reader
]
rows_html = "\n".join(rows)
except csv.Error as e:
raise ValueError(f"Error parsing CSV: {e}")
except StopIteration:
return "<p>No data found in CSV input.</p>"
rendered = f"""
<!DOCTYPE html>
<html>
<head>
<title>Query Results</title>
<style>{page_css}</style>
</head>
<body>
<table>
<thead>
<tr>{header_html}</tr>
</thead>
<tbody>
{rows_html}
</tbody>
</table>
<script>{page_script}</script>
</body>
</html>
"""
return rendered
def main():
try:
csv_data = sys.stdin.read()
html_output = csv_to_html(csv_data)
if html_output:
print(html_output)
except UnicodeDecodeError as e:
raise ValueError(
f"Error decoding input: {e}. Ensure your CSV data is encoded correctly (e.g., UTF-8)."
)
except Exception as e:
raise Exception(f"An unexpected error occurred: {e}")
if __name__ == "__main__":
main()
-939
View File
@@ -1,939 +0,0 @@
#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.14"
# dependencies = [
# "bottle>=0.13.0",
# ]
# ///
import argparse
import functools
import os
import re
import socket
import sys
import threading
import time
import webbrowser
from pathlib import Path
from datetime import datetime
from typing import TypedDict
import bottle
IMAGE_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".webp", ".webm"}
class ImageGroup(TypedDict):
"""Type definition for grouped images."""
base: str
original: Path
edits: list[Path]
def find_available_port() -> int:
"""Find an available random port."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
s.listen(1)
port = s.getsockname()[1]
return port
def is_image(path: str | Path) -> bool:
"""Check if file is a supported image."""
return Path(path).suffix.lower() in IMAGE_EXTENSIONS
def discover_images(paths: list[str]) -> list[Path]:
"""Discover all images from mixed list of files and folders."""
images = []
for path_str in paths:
path = Path(path_str)
if path.is_file() and is_image(path):
images.append(path.resolve())
elif path.is_dir():
for item in path.iterdir():
if item.is_file() and is_image(item):
images.append(item.resolve())
return sorted(set(images))
def group_images(images: list[Path]) -> list[ImageGroup]:
"""
Group images by original + edits.
Original: image.jpg
Edits: image-edit1.jpg, image-edit2.jpg, etc.
Returns list of dicts: {original: Path, edits: [Path, ...]}
"""
groups = {}
for img in images:
stem = img.stem
# Check if this is an edit
match = re.match(r"^(.+?)-edit", stem)
if match:
base = match.group(1)
if base not in groups:
groups[base] = {"original": None, "edits": []}
groups[base]["edits"].append(img)
else:
# This is an original
if stem not in groups:
groups[stem] = {"original": None, "edits": []}
groups[stem]["original"] = img
# Filter out groups without originals, convert to list
result = []
for base, group in groups.items():
if group["original"]:
result.append(ImageGroup(base=base, original=group["original"], edits=sorted(group["edits"])))
return sorted(result, key=lambda g: g["original"].name)
def safe_image_path(filepath: str) -> Path:
"""Verify image path is safe (prevent directory traversal)."""
# Search for the image in all_images by filename
for img in all_images:
if img.name == filepath:
return img
raise ValueError(f"Unauthorized path: {filepath}")
def hardlink_images(src_paths: list[Path | str], dest_dir: str | Path) -> None:
"""Hard link multiple images to destination directory."""
dest_path = Path(dest_dir)
make_dirs(dest_path)
for src in src_paths:
src = Path(src)
dest = dest_path / src.name
try:
# Remove existing file if present
if dest.exists():
dest.unlink()
os.link(src, dest)
except Exception as e:
print(f"Error hard linking {src} to {dest}: {e}", file=sys.stderr)
raise
# Global state
app = bottle.Bottle()
all_images = set()
image_groups = []
current_group_idx = 0
last_heartbeat = datetime.now()
should_exit = False
picks_dir = Path()
flags = []
flag_dirs = {}
HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Image Culler</title>
<script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
<style>
* {
margin: 0;
padding: 0;
box-sizing: border-box;
}
body {
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
background: #1e1e1e;
color: #e0e0e0;
overflow: hidden;
height: 100vh;
}
.container {
display: flex;
height: calc(100vh - 46px);
gap: 20px;
padding: 20px;
}
.panel {
display: flex;
flex-direction: column;
gap: 15px;
}
.left-panel {
flex: 0 0 40%;
}
.right-panel {
flex: 0 0 60%;
overflow-y: auto;
}
.image-display {
width: 100%;
height: 100%;
background: #2a2a2a;
border-radius: 8px;
display: flex;
align-items: center;
justify-content: center;
overflow: hidden;
border: 2px solid transparent;
transition: border-color 0.2s;
}
.image-display:hover {
border-color: #4a9eff;
}
.image-display.picked {
border-color: #4ade80;
background: rgba(74, 222, 128, 0.1);
}
.image-display img {
max-width: 100%;
max-height: 100%;
object-fit: contain;
}
.image-info {
background: #2a2a2a;
padding: 12px;
border-radius: 8px;
font-size: 14px;
}
dialog {
background: #2a2a2a;
color: #e0e0e0;
border: 1px solid #404040;
border-radius: 8px;
padding: 20px;
max-width: 400px;
}
dialog::backdrop {
background: rgba(0, 0, 0, 0.5);
}
.grid {
display: grid;
grid-template-columns: repeat(2, 1fr);
gap: 10px;
}
.grid[data-image-count="1"] {
grid-template-columns: 1fr;
}
.grid-item {
position: relative;
background: #2a2a2a;
border-radius: 8px;
overflow: hidden;
cursor: pointer;
border: 2px solid transparent;
transition: all 0.2s;
}
.grid-item:hover {
border-color: #4a9eff;
}
.grid-item img {
width: 100%;
height: auto;
display: block;
object-fit: contain;
}
.grid[data-image-count="1"] .grid-item img {
max-height: 90vh;
}
.grid-item.picked {
border-color: #4ade80;
background: rgba(74, 222, 128, 0.1);
}
.flag-pill {
background: #404040;
color: #e0e0e0;
border: 1px solid #555;
}
.flag-pill:hover {
background: #505050;
border-color: #666;
}
.flag-pill-active {
background: #4a9eff;
color: #000;
border: 1px solid #4a9eff;
font-weight: 500;
}
.flag-pill-active:hover {
background: #6ab0ff;
border-color: #6ab0ff;
}
</style>
</head>
<body>
<div x-data="app()" x-init="init()" @keydown.window="handleKeydown($event)">
<div style="background: #2a2a2a; padding: 12px 20px; border-bottom: 1px solid #404040;">
<div style="font-size: 16px; font-weight: 500; color: #e0e0e0; word-break: break-all;">
<span x-text="currentGroup ? currentGroup.original.name : 'No images'"></span>
</div>
</div>
<div class="container">
<!-- Left Panel: Original Image -->
<div class="panel left-panel">
<div class="image-display" :class="{ picked: originalPicked }" style="position: relative;">
<template x-if="currentGroup">
<img :src="`/image/${currentGroup.original.name}`" :alt="currentGroup.original.name">
<template x-if="imageFlagAssignments[currentGroup.original.name]">
<div style="position: absolute; bottom: 12px; left: 12px; background: #4a9eff; color: #000; padding: 4px 8px; border-radius: 4px; font-size: 12px; font-weight: bold;">
<span x-text="imageFlagAssignments[currentGroup.original.name]"></span>
</div>
</template>
</template>
<template x-if="!currentGroup">
<div style="text-align: center; color: #888;">No images</div>
</template>
</div>
<div class="image-info">
<template x-if="currentGroup">
<div>
<div style="display: flex; align-items: center; gap: 8px;">
<strong>Position:</strong>
<input
type="number"
:value="currentGroupIdx + 1"
@keyup.enter="jumpToIndex($event)"
@change="jumpToIndex($event)"
min="1"
:max="imageGroups.length"
style="width: 90px; background: #404040; color: #e0e0e0; border: 1px solid #555; padding: 4px 8px; border-radius: 4px; font-size: 14px;"
>
<span style="color: #888;">/ <span x-text="imageGroups.length"></span></span>
</div>
<template x-if="availableFlags.length > 0">
<div style="margin-top: 10px;">
<strong style="font-size: 13px;">Flags:</strong>
<div style="display: flex; flex-wrap: wrap; gap: 6px; margin-top: 6px;">
<template x-for="(flag, idx) in availableFlags" :key="idx">
<div
@click="handleFlagAssignment(flag)"
:class="imageFlagAssignments[currentGroup.original.name] === flag ? 'flag-pill-active' : 'flag-pill'"
style="cursor: pointer; padding: 4px 8px; border-radius: 4px; font-size: 12px; white-space: nowrap; transition: all 0.2s;"
>
<span x-text="`${flag} (${idx + 1})`"></span>
</div>
</template>
</div>
</div>
</template>
</div>
</template>
</div>
</div>
<!-- Right Panel: Edit Variations Grid -->
<div class="panel right-panel">
<template x-if="currentGroup && currentGroup.edits.length > 0">
<div class="grid" :data-image-count="currentGroup.edits.length">
<template x-for="(edit, idx) in currentGroup.edits" :key="idx">
<div
class="grid-item"
:class="{ picked: pickedEdits.includes(idx) }"
@click="toggleEditPick(idx)"
>
<img :src="`/image/${edit.name}`" :alt="`Edit ${idx + 1}`">
</div>
</template>
</div>
</template>
<template x-if="!currentGroup || currentGroup.edits.length === 0">
<div style="color: #888; text-align: center; padding: 40px;">No variations</div>
</template>
</div>
</div>
<!-- Help Dialog -->
<dialog id="helpDialog">
<h2>Keyboard Shortcuts</h2>
<div style="margin-top: 15px; font-size: 14px; line-height: 1.8;">
<div style="margin-bottom: 10px;">
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">P</span>
Toggle pick (all edits or hover one)
</div>
<div style="margin-bottom: 10px;">
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">←→</span>
Navigate any image
</div>
<div style="margin-bottom: 10px;">
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">K</span>
Skip to next pending
</div>
<div style="margin-top: 15px; font-size: 13px; color: #aaa;">
<strong>Flags:</strong>
</div>
<div>
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">1-9</span>
Assign/unassign to flag
</div>
<div>
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">/</span>
Toggle reject flag
</div>
</div>
<div style="margin-top: 20px; text-align: right;">
<button @click="closeHelp()" style="background: #404040; color: #e0e0e0; border: none; padding: 8px 16px; border-radius: 4px; cursor: pointer;">Close</button>
</div>
</dialog>
</div>
<script>
function app() {
return {
imageGroups: [],
currentGroupIdx: 0,
currentGroup: null,
pickedEdits: [],
originalPicked: false,
groupStatus: {},
availableFlags: [],
imageFlagAssignments: {},
async init() {
// Start heartbeat
this.startHeartbeat();
// Track mouse position globally for hover detection
document.addEventListener('mousemove', (e) => {
window.lastMouseEvent = e;
});
// Notify server when tab/window is closed
window.addEventListener('beforeunload', () => {
fetch('/api/shutdown', { method: 'POST' }).catch(() => {});
});
// Load images
await this.loadImages();
// Load first group
this.loadGroup(0);
},
async loadImages() {
try {
const resp = await fetch('/api/images');
const data = await resp.json();
this.imageGroups = data.groups;
this.availableFlags = data.flags_list || [];
this.groupStatus = {};
this.imageFlagAssignments = data.flag_assignments || {};
this.imageGroups.forEach((_, idx) => {
this.groupStatus[idx] = 'pending';
});
} catch (err) {
console.error('Error loading images:', err.message);
}
},
loadGroup(idx) {
if (idx < 0 || idx >= this.imageGroups.length) return;
this.currentGroupIdx = idx;
this.currentGroup = this.imageGroups[idx];
this.pickedEdits = [];
this.originalPicked = false;
},
getGroupStatus(idx) {
return this.groupStatus[idx] || 'pending';
},
jumpToIndex(event) {
const value = parseInt(event.target.value);
if (!isNaN(value) && value >= 1 && value <= this.imageGroups.length) {
this.loadGroup(value - 1);
}
},
toggleEditPick(idx) {
if (this.pickedEdits.includes(idx)) {
this.pickedEdits = this.pickedEdits.filter(i => i !== idx);
} else {
this.pickedEdits.push(idx);
}
},
getHoveredElement() {
// Find the element currently under the mouse pointer
const mouseEvent = window.lastMouseEvent;
if (!mouseEvent) return null;
return document.elementFromPoint(mouseEvent.clientX, mouseEvent.clientY);
},
async handleKeydown(e) {
// Skip keybindings if focused on input/textarea
const target = e.target;
const isInput = target.tagName === 'INPUT' || target.tagName === 'TEXTAREA' || target.contentEditable === 'true';
if (isInput) {
return;
}
const key = e.key.toLowerCase();
if (key === 'p') {
e.preventDefault();
await this.handlePick();
} else if (key === '?') {
e.preventDefault();
this.toggleHelp();
} else if (key === '/') {
e.preventDefault();
await this.handleFlagAssignment('reject');
} else if (key === 'arrowleft') {
e.preventDefault();
let nextIdx = this.currentGroupIdx - 1;
if (nextIdx < 0) nextIdx = this.imageGroups.length - 1;
this.loadGroup(nextIdx);
} else if (key === 'arrowright') {
e.preventDefault();
let nextIdx = this.currentGroupIdx + 1;
if (nextIdx >= this.imageGroups.length) nextIdx = 0;
this.loadGroup(nextIdx);
} else if (key === 'k') {
e.preventDefault();
for (let i = this.currentGroupIdx + 1; i < this.imageGroups.length; i++) {
if (this.getGroupStatus(i) === 'pending') {
this.loadGroup(i);
return;
}
}
} else if (key >= '1' && key <= '9') {
e.preventDefault();
const flagIdx = parseInt(key) - 1;
if (flagIdx < this.availableFlags.length) {
const flagName = this.availableFlags[flagIdx];
await this.handleFlagAssignment(flagName);
}
}
},
async handlePick() {
if (!this.currentGroup) return;
const currentStatus = this.getGroupStatus(this.currentGroupIdx);
// If already picked, unpick it
if (currentStatus === 'picked') {
try {
const resp = await fetch('/api/unpick', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: this.currentGroup.original.name
})
});
if (resp.ok) {
this.groupStatus[this.currentGroupIdx] = 'pending';
this.originalPicked = false;
this.pickedEdits = [];
}
} catch (err) {
console.error('Error unpicking:', err);
}
return;
}
const hoveredEl = this.getHoveredElement();
const isHoveringOriginal = hoveredEl?.closest('.image-display:not(.right-panel *)');
// If hovering over original, pick just the original
if (isHoveringOriginal) {
try {
const resp = await fetch('/api/pick', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: this.currentGroup.original.name,
edits: [],
includeOriginal: true
})
});
if (resp.ok) {
this.originalPicked = true;
}
} catch (err) {
console.error('Error picking original:', err);
}
return;
}
// If hovering on edit, pick just that edit
const hoveredGridItem = hoveredEl?.closest('.grid-item');
if (hoveredGridItem && this.currentGroup.edits.length > 0) {
const idx = Array.from(hoveredGridItem.parentElement.querySelectorAll('.grid-item')).indexOf(hoveredGridItem);
if (idx !== -1) {
const editName = this.currentGroup.edits[idx].name;
try {
const resp = await fetch('/api/pick', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: this.currentGroup.original.name,
edits: [editName],
includeOriginal: false
})
});
if (resp.ok) {
this.pickedEdits.push(idx);
}
} catch (err) {
console.error('Error picking edit:', err);
}
return;
}
}
// If not hovering, pick all selected edits and advance
let editsToPick = this.pickedEdits.map(i => this.currentGroup.edits[i].name);
// If no edits selected but there's exactly 1 edit, pick it automatically
if (editsToPick.length === 0 && this.currentGroup.edits.length === 1) {
editsToPick = [this.currentGroup.edits[0].name];
}
// Only proceed if there are edits to pick
if (editsToPick.length === 0) return;
try {
const resp = await fetch('/api/pick', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: this.currentGroup.original.name,
edits: editsToPick,
includeOriginal: false
})
});
if (resp.ok) {
this.groupStatus[this.currentGroupIdx] = 'picked';
this.loadGroup(this.currentGroupIdx + 1);
}
} catch (err) {
console.error('Error picking edits:', err);
}
},
startHeartbeat() {
setInterval(async () => {
try {
await fetch('/api/heartbeat', { method: 'POST' });
} catch (err) {
console.error('Heartbeat failed:', err);
}
}, 1000);
},
toggleHelp() {
const dialog = document.getElementById('helpDialog');
if (dialog.open) {
dialog.close();
} else {
dialog.showModal();
}
},
closeHelp() {
document.getElementById('helpDialog').close();
},
async handleFlagAssignment(flagName) {
if (!this.currentGroup) return;
const imageKey = this.currentGroup.original.name;
// Toggle: if already assigned to this flag, unassign
if (this.imageFlagAssignments[imageKey] === flagName) {
try {
await fetch('/api/unassign-flag', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: imageKey,
flag: flagName
})
});
delete this.imageFlagAssignments[imageKey];
} catch (err) {
console.error('Error unassigning flag:', err);
}
} else {
// Assign to new flag
try {
const resp = await fetch('/api/assign-flag', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({
original: imageKey,
flag: flagName
})
});
if (resp.ok) {
this.imageFlagAssignments[imageKey] = flagName;
}
} catch (err) {
console.error('Error assigning flag:', err);
}
}
}
};
}
</script>
</body>
</html>
"""
@app.get("/")
def index() -> str:
return HTML_TEMPLATE
@app.get("/api/images")
def api_images():
"""Return grouped images and current flag assignments."""
global image_groups, flags, flag_dirs
result = []
for group in image_groups:
result.append(
{
"base": group["base"],
"original": {"name": group["original"].name, "path": str(group["original"])},
"edits": [{"name": e.name, "path": str(e)} for e in group["edits"]],
}
)
# Build flag assignments by checking which flag directories contain which original images
flag_assignments = {}
for flag, flag_dir in flag_dirs.items():
if flag_dir.exists():
for item in flag_dir.iterdir():
if item.is_file():
# Map the file name to its flag
flag_assignments[item.name] = flag
return {"groups": result, "flags_list": flags, "flag_assignments": flag_assignments}
@app.get("/image/<filename>")
def serve_image(filename: str):
"""Serve an image file safely."""
try:
target = safe_image_path(filename)
return bottle.static_file(target.name, root=str(target.parent), mimetype="image/jpeg")
except ValueError:
bottle.response.status = 403
return {"error": "Unauthorized path"}
@app.post("/api/pick")
def api_pick():
"""Hard link picked images to _picks directory."""
global image_groups, current_group_idx, picks_dir
# Create directory before writing
make_dirs(picks_dir)
data = bottle.request.json
original_name = data.get("original")
edits_names = data.get("edits", [])
include_original = data.get("includeOriginal", False)
files_to_pick = []
# Add original if requested
if include_original:
for img in all_images:
if img.name == original_name:
files_to_pick.append(img)
break
# Add edits by finding them in all_images
for edit_name in edits_names:
for img in all_images:
if img.name == edit_name:
files_to_pick.append(img)
break
try:
hardlink_images(files_to_pick, picks_dir)
return {"status": "ok"}
except Exception as e:
bottle.response.status = 500
return {"error": str(e)}
@app.post("/api/unpick")
def api_unpick():
"""Remove picked images from _picks directory."""
global picks_dir
data = bottle.request.json
original_name = data.get("original")
try:
# Find and remove the original and any edits from picks directory
if picks_dir.exists():
for item in picks_dir.iterdir():
if item.is_file() and (item.name == original_name or item.stem.startswith(original_name.rsplit(".", 1)[0] + "-edit")):
item.unlink()
return {"status": "ok"}
except Exception as e:
bottle.response.status = 500
return {"error": str(e)}
@app.post("/api/assign-flag")
def api_assign_flag():
"""Assign an image to a flag."""
global flag_dirs, all_images
data = bottle.request.json
original_name = data.get("original")
flag_name = data.get("flag")
# Find original file
original_path = None
for img in all_images:
if img.name == original_name:
original_path = img
break
if not original_path or flag_name not in flag_dirs:
bottle.response.status = 400
return {"error": "Invalid image or flag"}
# Create directory before writing
make_dirs(flag_dirs[flag_name])
try:
hardlink_images([original_path], flag_dirs[flag_name])
return {"status": "ok"}
except Exception as e:
bottle.response.status = 500
return {"error": str(e)}
@app.post("/api/unassign-flag")
def api_unassign_flag():
"""Remove an image from a flag."""
global flag_dirs
data = bottle.request.json
original_name = data.get("original")
flag_name = data.get("flag")
if flag_name not in flag_dirs:
bottle.response.status = 400
return {"error": "Invalid flag"}
try:
target_file = flag_dirs[flag_name] / original_name
if target_file.exists():
target_file.unlink()
return {"status": "ok"}
except Exception as e:
bottle.response.status = 500
return {"error": str(e)}
@app.post("/api/heartbeat")
def api_heartbeat():
"""Update last heartbeat timestamp."""
global last_heartbeat
last_heartbeat = datetime.now()
return {"status": "ok"}
@app.post("/api/shutdown")
def api_shutdown():
"""Shutdown the server."""
global should_exit
should_exit = True
return {"status": "ok"}
@functools.cache
def make_dirs(path: Path) -> None:
"""Create directory if it doesn't exist."""
path.mkdir(parents=True, exist_ok=True)
def heartbeat_monitor() -> None:
"""Monitor for shutdown signal."""
global should_exit
while not should_exit:
time.sleep(0.5)
print("\nShutting down server.", file=sys.stderr)
os._exit(0)
def main() -> None:
global all_images, image_groups, last_heartbeat, picks_dir, flags, flag_dirs
parser = argparse.ArgumentParser(description="Image culling app for AI image2image transforms")
parser.add_argument("paths", nargs="+", help="Image files or folders to cull")
parser.add_argument("--picks-dir", default=None, help="Directory for picked images (default: $cwd/_picks)")
parser.add_argument("--rejects-dir", default=None, help="Directory for rejected images (default: $cwd/_rejects)")
parser.add_argument("--flags", default=None, help="Comma-separated flag names (e.g., 'a,b,c,d')")
parser.add_argument("--flags-dir", default=None, help="Directory for flag subdirectories (default: $cwd)")
args = parser.parse_args()
# Set directories
cwd = Path.cwd()
picks_dir = Path(args.picks_dir or (cwd / "_picks"))
rejects_dir = Path(args.rejects_dir or (cwd / "_rejects"))
flags_dir = Path(args.flags_dir or cwd)
# Set up flags
if args.flags:
flags = [f.strip() for f in args.flags.split(",")]
for flag in flags:
flag_dir = flags_dir / f"_picks_flag_{flag}"
flag_dirs[flag] = flag_dir
# Always add reject as a flag at the end
flags.append("reject")
flag_dirs["reject"] = rejects_dir
# Discover images
images = discover_images(args.paths)
all_images = set(images)
if not images:
print("No images found in provided paths.", file=sys.stderr)
sys.exit(1)
# Group images
image_groups = group_images(images)
if not image_groups:
print("No original images found (looking for files matching naming convention).", file=sys.stderr)
sys.exit(1)
print(f"Found {len(image_groups)} original image(s) with variations", file=sys.stderr)
# Find available port
port = find_available_port()
url = f"http://127.0.0.1:{port}"
# Start heartbeat monitor
monitor_thread = threading.Thread(target=heartbeat_monitor, daemon=True)
monitor_thread.start()
# Open browser
print(f"Opening browser at {url}", file=sys.stderr)
webbrowser.open(url)
# Start Flask server
last_heartbeat = datetime.now()
bottle.run(app, host="127.0.0.1", port=port, quiet=True)
if __name__ == "__main__":
main()
-194
View File
@@ -1,194 +0,0 @@
#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["playwright", "httpx"]
# ///
import argparse
import json
import logging
import re
import sys
from pathlib import Path
from typing import Optional
from urllib.parse import urljoin, urlparse
import httpx
from playwright.sync_api import sync_playwright
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger(__name__)
def str_to_bool(value: str) -> bool:
"""Convert string to boolean."""
if isinstance(value, bool):
return value
return value.lower() in ("true", "1", "yes", "on")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Download JavaScript files and reconstruct sources from source maps",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("url", help="URL of the page to scrape")
parser.add_argument(
"--output-dir",
"-o",
type=Path,
default=Path("output"),
help="Output directory for reconstructed sources",
)
parser.add_argument(
"--headless",
type=str_to_bool,
default=True,
help="Run browser in headless mode",
)
return parser.parse_args()
def get_script_urls(page_url: str, headless: bool = True) -> list[str]:
"""Extract all script URLs from a page using Playwright."""
script_urls = []
with sync_playwright() as playwright:
browser = playwright.chromium.launch(headless=headless)
page = browser.new_page()
def handle_response(response):
if response.request.resource_type == "script":
script_urls.append(response.url)
logger.debug(f"Detected script: {response.url}")
page.on("response", handle_response)
logger.info("Navigating to page...")
page.goto(url=page_url)
logger.info("Waiting for DOM ready...")
page.wait_for_load_state(state="domcontentloaded")
logger.info("Waiting for network idle (30s timeout)...")
try:
page.wait_for_load_state(state="networkidle", timeout=30000)
except Exception as e:
logger.warning(f"Network idle timeout: {e}")
browser.close()
return script_urls
def download_file(url: str, timeout: int = 30) -> bytes:
"""Download a file from a URL."""
try:
with httpx.Client(timeout=timeout, follow_redirects=True) as client:
response = client.get(url=url)
response.raise_for_status()
return response.content
except httpx.HTTPError as e:
logger.error(f"Failed to download {url}: {e}")
raise
def get_sourcemap_url(script_content: str, script_url: str) -> Optional[str]:
"""Extract source map URL from script content."""
match = re.search(r"//# sourceMappingURL=(.+?)(?:\n|$)", script_content)
if not match:
return None
sourcemap_ref = match.group(1).strip()
if sourcemap_ref.startswith("http"):
return sourcemap_ref
return urljoin(base=script_url, url=sourcemap_ref)
def decode_sourcemap(sourcemap_data: dict) -> dict:
"""Extract sources and content from source map."""
sources = sourcemap_data.get("sources", [])
sources_content = sourcemap_data.get("sourcesContent", [])
return {
"sources": sources,
"sources_content": sources_content,
}
def reconstruct_sources(
script_content: str,
sourcemap_data: dict,
output_dir: Path,
script_name: str,
) -> None:
"""Reconstruct original sources from source map."""
sources = sourcemap_data.get("sources", [])
sources_content = sourcemap_data.get("sourcesContent", [])
logger.info(f"Source map has {len(sources)} source files")
if sources:
logger.debug(f"First few sources: {sources[:3]}")
if not sources_content or all(c is None for c in sources_content):
logger.error(f"No source content found in source map for {script_name}")
sys.exit(1)
for i, source_file in enumerate(sources):
if i < len(sources_content) and sources_content[i]:
source_path = output_dir / source_file
source_path.parent.mkdir(parents=True, exist_ok=True)
with open(source_path, mode="w", encoding="utf-8") as f:
f.write(sources_content[i])
logger.info(f"Reconstructed {source_file}")
def main() -> None:
args = parse_args()
args.output_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Fetching page: {args.url}")
script_urls = get_script_urls(page_url=args.url, headless=args.headless)
logger.info(f"Found {len(script_urls)} script(s)")
if not script_urls:
logger.error("No scripts found on page")
sys.exit(1)
for script_url in script_urls:
logger.info(f"Processing {script_url}")
try:
script_content = download_file(url=script_url).decode(encoding="utf-8")
except Exception as e:
logger.error(f"Failed to get script {script_url}: {e}")
continue
sourcemap_url = get_sourcemap_url(script_content=script_content, script_url=script_url)
if not sourcemap_url:
logger.error(f"No source map found for {script_url}")
sys.exit(1)
logger.info(f"Downloading source map: {sourcemap_url}")
try:
sourcemap_content = download_file(url=sourcemap_url).decode(encoding="utf-8")
sourcemap_data = json.loads(sourcemap_content)
except Exception as e:
logger.error(f"Failed to download/parse source map: {e}")
sys.exit(1)
script_name = urlparse(script_url).path.split("/")[-1]
reconstruct_sources(
script_content=script_content,
sourcemap_data=sourcemap_data,
output_dir=args.output_dir,
script_name=script_name,
)
logger.info(f"All sources reconstructed to {args.output_dir}")
if __name__ == "__main__":
main()
-40
View File
@@ -1,40 +0,0 @@
import sys
import re
import math
import datetime
def parse_duration(text: str) -> datetime.timedelta:
re_duration = re.compile(r'(\d+(?:\.\d+)?(?:sec|min|hour|hr|h|m|s))')
re_digits = re.compile(r'([\d.]+)')
time_scales = {
's': 1,
'sec': 1,
'm': 60,
'min': 60,
'h': 3600,
'hr': 3600,
}
parts = re_duration.findall(text)
if parts:
total_secs = 0
for it in parts:
if m := re_digits.search(it):
_, end = m.span()
scale = it[end:]
total_secs += math.ceil(float(m.group(1)) * time_scales[scale])
return datetime.timedelta(seconds=total_secs)
parts = [float(it) for it in text.split(':')]
if len(parts) == 1:
return datetime.timedelta(seconds=parts[0])
if len(parts) == 2:
return datetime.timedelta(minutes=parts[0], seconds=parts[1])
if len(parts) == 3:
return datetime.timedelta(hours=parts[0], minutes=parts[1], seconds=parts[2])
if __name__ == '__main__':
parsed = parse_duration(sys.argv[1])
seconds = int(parsed.total_seconds())
-230
View File
@@ -1,230 +0,0 @@
#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.13"
# ///
import argparse
import logging
import re
import subprocess
import sys
from pathlib import Path
logger = logging.getLogger(__name__)
KNOWN_PASSWORDS = """
torgo
Yuzuki
oron.com
MonikGonPs
koth
drdoom-psuzy
psuzy
suzy
idle
dan4260
""".strip().splitlines(keepends=False)
class CorruptArchiveError(Exception):
pass
def find_archive_files(root_dir: Path) -> list[Path]:
extractables = []
globs = ["*.rar", "*.zip*", "*.7z*"]
for glob in globs:
for f in root_dir.rglob(glob):
extractables.append(f)
extractables = [f for f in extractables if is_extractable(f)]
return sorted(extractables, key=str)
def is_extractable(archive: Path) -> bool:
if ".corrupt" in archive.name:
return False
if re.search(r"\.zip|\.7z", archive.name):
zip_part = re.search(r"(\d+)$", archive.suffix)
if zip_part and int(zip_part.group(1)) == 1:
return True
if not zip_part:
return True
if archive.name.endswith(".rar"):
rar_part = re.search(r"\.part(\d+)", archive.name)
if not rar_part:
return True
if rar_part and int(rar_part.group(1)) == 1:
if archive.with_name(archive.name.replace("part1", "part2")).is_file():
return True
return False
def extract_archive_rar(archive: Path, output_dir: Path, password: str | None = None) -> None:
args = [
"rar",
"e",
*([f"-p{password}"] if password else []),
f"{archive.name}",
f"{output_dir}",
]
try:
subprocess.run(args, cwd=archive.parent, capture_output=True, text=True, check=True)
except subprocess.CalledProcessError as e:
output: str = e.stderr or e.stdout
if "corrupt file or wrong password" in output.lower():
raise LookupError("corrupt or invalid password")
elif e.returncode == 3:
raise FileNotFoundError("missing parts")
elif e.returncode == 10:
raise ValueError("not a rar file")
elif e.returncode == 11:
raise LookupError("invalid password")
raise
def extract_archive_7z(archive: Path, output_dir: Path, password: str | None = None):
args = [
"7z",
"e",
*(["-p", password] if password else []),
f"-o{output_dir}",
"-y",
f"{archive}",
]
try:
subprocess.run(args, capture_output=True, text=True, check=True)
except subprocess.CalledProcessError as e:
output: str = e.stderr or e.stdout
if "missing volume" in output.lower():
raise CorruptArchiveError("missing parts")
if "wrong password" in output.lower():
raise LookupError("invalid password")
raise
def extract_archive_zip(archive: Path, output_dir: Path, password: str | None = None):
args = [
"unzip",
"-o", # overwrite files without prompting
*(["-P", password] if password else []),
f"{archive}",
"-d",
f"{output_dir}",
]
try:
subprocess.run(args, capture_output=True, text=True, check=True)
except subprocess.CalledProcessError as e:
output: str = e.stderr or e.stdout
if "zipfile corrupt" in output.lower():
raise CorruptArchiveError("corrupt zip file")
if "cannot find zipfile" in output.lower():
raise FileNotFoundError("zip file not found")
if "incorrect password" in output.lower() or "password required" in output.lower():
raise LookupError("invalid password")
raise
def extract_archive(archive_path: Path, output_dir: Path, passwords: list[str]) -> None:
for p in ["", *passwords]:
try:
fn: callable
if ".rar" in archive_path.name:
fn = extract_archive_rar
elif ".zip" == archive_path.suffix:
fn = extract_archive_zip
else:
fn = extract_archive_7z
fn(archive_path, output_dir, p)
except subprocess.CalledProcessError as e:
logger.error(f"error extracting {archive_path}: {e.stderr}")
raise
except LookupError as e:
logger.info(f"invalid password: {p!r}. {e}")
continue
except Exception as e:
logger.error(f"failed to extract: {e}")
raise
def delete_archive(archive: Path):
logger.info(f"deleting: {archive.name}")
archive.unlink()
stem = archive.stem.rsplit(".", maxsplit=1)[0]
for f in archive.parent.rglob(f"*{archive.suffix}"):
if archive.stem.startswith(stem):
logger.info(f"deleting: {f.name}")
f.unlink()
def unlink_if_empty(dir_path: Path) -> None:
"""Unlink the directory if it is empty."""
try:
if not any(dir_path.iterdir()):
dir_path.rmdir()
logger.info(f"deleted empty directory: {dir_path}")
except OSError as e:
logger.error(f"error deleting directory: {e}")
raise
def parse_args():
arger = argparse.ArgumentParser()
arger.add_argument("archive_paths", type=lambda v: Path(v.strip()), nargs="+", help="List of archive paths")
arger.add_argument("--cwd", type=Path, default=Path.cwd(), help="working directory")
arger.add_argument("--find", default=False, help="find archives and print to stdout", action="store_true")
arger.add_argument("--clean", default=False, help="delete file(s) after successful extraction", action="store_true")
if len(sys.argv) < 2:
arger.print_help()
exit(1)
return arger.parse_args()
def main():
logging.basicConfig(level=logging.INFO)
args = parse_args()
cwd: Path = args.cwd
if args.find:
archives = find_archive_files(cwd)
for f in archives:
print(str(f.resolve()))
return
archives = args.archive_paths
failed = []
for i, f in enumerate(archives, start=1):
f: Path
progress = f"{i:02}/{len(archives):02}"
logger.info(f"extracting {progress}: {f.relative_to(cwd)}")
output_dir = f.parent / f.stem.strip()
try:
output_dir.mkdir(exist_ok=True, parents=True)
extract_archive(f, output_dir=output_dir, passwords=KNOWN_PASSWORDS)
if args.clean:
delete_archive(f)
except CorruptArchiveError as e:
if ".corrupt" not in f.name:
f.rename(f.with_stem(f"{f.stem}.corrupt"))
except (ValueError, FileNotFoundError) as e:
logger.error(str(e))
failed.append(f)
finally:
unlink_if_empty(output_dir)
if failed:
logger.info(f"failed to extract {len(failed)} archives")
for f in failed:
logger.info(f"{f.relative_to(cwd)}")
if __name__ == "__main__":
main()
-81
View File
@@ -1,81 +0,0 @@
#!/usr/bin/env python3.9
import argparse
import json
import logging
import subprocess
import sys
from pathlib import Path
import typing
def parse_args(argv: list[str]):
arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
arger.add_argument("video_path", type=Path, help="Video path")
arger.add_argument("save_path", nargs="?", type=Path, help="Output path")
if len(argv) == 0:
arger.print_help()
exit(1)
return arger.parse_args(argv)
def ffprobe(video_path: Path) -> dict:
proc = subprocess.run(
[
"ffprobe",
"-v",
"quiet",
"-print_format",
"json",
"-show_format",
"-show_streams",
str(video_path),
],
check=True,
capture_output=True,
)
return json.loads(proc.stdout)
def clean_video(video_path: Path, save_path: Path) -> Path:
media_info = ffprobe(video_path)
# check if the first video stream has hevc codec
if media_info["streams"][0]["codec_name"] == "hevc":
logging.debug(f"Video codec is hevc, adding hvc1 tag")
extra_args = ["-tag:v", "hvc1"]
else:
extra_args = []
# fmt: off
args = [
'ffmpeg',
'-i',
str(video_path),
'-c:v', 'copy',
# '-movflags', '+faststart',
'-map_metadata', '-1',
*extra_args,
'-c:a', 'copy',
'-y',
str(save_path),
]
# fmt: on
logging.info(f"calling ffmpeg with {args=}")
subprocess.run(args, check=True)
return save_path
def main():
logging.basicConfig(level=logging.INFO)
args = parse_args(sys.argv[1:])
save_path: typing.Optional[Path] = args.save_path
if save_path and save_path.resolve().is_dir():
save_path = save_path / args.video_path.name
if save_path is None:
save_path = args.video_path.with_suffix(".clean.mp4")
clean_video(video_path=args.video_path, save_path=save_path)
if __name__ == "__main__":
main()
-461
View File
@@ -1,461 +0,0 @@
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.10,<3.11"
# dependencies = [
# "torch==2.9.1+cu128",
# "torchvision==0.24.1+cu128",
# "torchaudio==2.9.1+cu128",
# "diffusers @ git+https://github.com/huggingface/diffusers.git",
# "transformers>=4.57.0",
# "accelerate>=1.0.0",
# "safetensors>=0.4.0",
# "huggingface_hub>=0.25.0",
# "numpy>=1.26.0",
# "pillow>=10.0.0",
# "fastapi>=0.115.0",
# "uvicorn>=0.29.0",
# ]
# [tool.uv]
# extra-index-url = ["https://download.pytorch.org/whl/cu128"]
# ///
import argparse
import base64
import binascii
import io
import os
import random
import threading
import time
from typing import Any
os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
import torch
from diffusers import Flux2KleinPipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse, Response
from pydantic import BaseModel, Field
from PIL import Image
def _resolve_model_size() -> str:
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--model", choices=["4b", "9b"], default="9b")
args, _ = parser.parse_known_args()
return args.model
def _checkpoint_for(model_size: str) -> str:
if model_size == "4b":
return "black-forest-labs/FLUX.2-klein-4B"
return "black-forest-labs/FLUX.2-klein-9B"
MODEL_SIZE = _resolve_model_size()
CHECKPOINT = _checkpoint_for(MODEL_SIZE)
USE_COMPILE = os.environ.get("FLUX2_COMPILE", "1") != "0"
WARMUP_RUNS = int(os.environ.get("FLUX2_WARMUP_RUNS", "2"))
WARMUP_STEPS = int(os.environ.get("FLUX2_WARMUP_STEPS", "4"))
WARMUP_PROMPT = "A high-resolution photo of a red fox in natural daylight"
MAX_SEED = (2**31) - 1
MAX_PIXELS = 4 * 1024 * 1024
DEFAULT_STEPS = 4
DEFAULT_GUIDANCE = 1.0
app = FastAPI(title="FLUX2 Klein Private Server")
PIPE: Flux2KleinPipeline | None = None
PIPE_LOCK = threading.Lock()
class GenerateRequest(BaseModel):
prompt: str = Field(..., min_length=1)
width: int = 1024
height: int = 1024
num_inference_steps: int = DEFAULT_STEPS
guidance_scale: float = DEFAULT_GUIDANCE
seed: int = 0
randomize_seed: bool = False
image_data_url: str | None = None
class CompatRequest(BaseModel):
model: str | None = None
prompt: str = Field(..., min_length=1)
imageDataUrl: str | None = None
size: str | None = None
width: int | None = None
height: int | None = None
num_inference_steps: int | None = None
guidance_scale: float | None = None
seed: int | None = None
randomize_seed: bool = False
response_format: str = "b64_json"
n: int = 1
negative_prompt: str | None = None
negativePrompt: str | None = None
def _ensure_cuda() -> None:
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required for this server")
def _compile_pipe(pipe: Flux2KleinPipeline) -> bool:
if not USE_COMPILE:
return False
print("Compiling transformer and VAE decoder...")
start = time.perf_counter()
pipe.transformer = torch.compile(
pipe.transformer,
mode="max-autotune",
dynamic=False,
fullgraph=True,
)
pipe.vae.decode = torch.compile(
pipe.vae.decode,
mode="max-autotune",
dynamic=False,
fullgraph=True,
)
print(f"Compile finished in {time.perf_counter() - start:.2f}s")
return True
def _warmup_pipe(pipe: Flux2KleinPipeline) -> None:
if WARMUP_RUNS <= 0:
return
print(f"Warmup started ({WARMUP_RUNS} runs)...")
start = time.perf_counter()
for run_idx in range(WARMUP_RUNS):
generator = torch.Generator(device="cuda").manual_seed(run_idx)
pipe(
prompt=WARMUP_PROMPT,
height=1024,
width=1024,
num_inference_steps=WARMUP_STEPS,
guidance_scale=1.0,
generator=generator,
).images[0]
torch.cuda.synchronize()
print(f"Warmup finished in {time.perf_counter() - start:.2f}s")
def _load_pipe(ckpt: str) -> Flux2KleinPipeline:
print(f"Loading checkpoint: {ckpt}")
pipe = Flux2KleinPipeline.from_pretrained(ckpt, torch_dtype=torch.bfloat16)
pipe = pipe.to("cuda")
pipe.transformer.fuse_qkv_projections()
pipe.vae.fuse_qkv_projections()
pipe.vae.to(memory_format=torch.channels_last)
pipe.transformer.set_attention_backend("_native_flash")
if _compile_pipe(pipe):
_warmup_pipe(pipe)
return pipe
def _normalize_dims(width: int, height: int) -> tuple[int, int]:
if width <= 0 or height <= 0:
raise HTTPException(status_code=422, detail="width and height must be > 0")
if width % 16 != 0 or height % 16 != 0:
raise HTTPException(
status_code=422, detail="width and height must be multiples of 16"
)
if width * height > MAX_PIXELS:
raise HTTPException(status_code=422, detail="max resolution is 4 megapixels")
return width, height
def _parse_size(
size: str | None, width: int | None, height: int | None
) -> tuple[int, int]:
if width and height:
return _normalize_dims(width, height)
if not size:
return _normalize_dims(1024, 1024)
parts = size.lower().split("x")
if len(parts) != 2:
raise HTTPException(status_code=422, detail="size must look like 1024x1024")
try:
parsed_w = int(parts[0])
parsed_h = int(parts[1])
except ValueError as exc:
raise HTTPException(
status_code=422, detail="size values must be integers"
) from exc
return _normalize_dims(parsed_w, parsed_h)
def _decode_image_data_url(image_data_url: str | None) -> Image.Image | None:
if not image_data_url:
return None
if "," not in image_data_url:
raise HTTPException(status_code=422, detail="imageDataUrl must be a data URL")
_, payload = image_data_url.split(",", 1)
try:
raw = base64.b64decode(payload)
except (ValueError, binascii.Error) as exc:
raise HTTPException(
status_code=422, detail="invalid base64 in imageDataUrl"
) from exc
with Image.open(io.BytesIO(raw)) as image:
return image.convert("RGB")
def _generate_png_bytes(
prompt: str,
width: int,
height: int,
num_inference_steps: int,
guidance_scale: float,
seed: int,
image: Image.Image | None,
) -> tuple[bytes, int]:
if PIPE is None:
raise HTTPException(status_code=503, detail="model not loaded yet")
generator = torch.Generator(device="cuda").manual_seed(seed)
kwargs: dict[str, Any] = {
"prompt": prompt,
"height": height,
"width": width,
"num_inference_steps": num_inference_steps,
"guidance_scale": guidance_scale,
"generator": generator,
}
with PIPE_LOCK:
try:
if image is not None:
try:
output = PIPE(image=image, **kwargs)
except TypeError:
output = PIPE(images=[image], **kwargs)
else:
output = PIPE(**kwargs)
except Exception as exc:
raise HTTPException(
status_code=500, detail=f"generation failed: {exc}"
) from exc
result_image = output.images[0]
buf = io.BytesIO()
result_image.save(buf, format="PNG")
return buf.getvalue(), seed
@app.on_event("startup")
def _startup() -> None:
global PIPE
_ensure_cuda()
PIPE = _load_pipe(CHECKPOINT)
print("Server ready")
@app.get("/", response_class=HTMLResponse)
def index() -> str:
return """<!doctype html>
<html lang=\"en\">
<head>
<meta charset=\"utf-8\" />
<meta name=\"viewport\" content=\"width=device-width, initial-scale=1\" />
<title>FLUX.2 Klein Playground</title>
<style>
:root { --bg:#f4f0e8; --card:#fff9ee; --text:#1f1f1f; --accent:#125b50; --muted:#6b6b6b; }
* { box-sizing: border-box; }
body { margin:0; font-family: ui-sans-serif, -apple-system, Segoe UI, sans-serif; background: radial-gradient(circle at top left,#fff6df,var(--bg)); color:var(--text); }
.wrap { max-width: 980px; margin: 1.5rem auto; padding: 1rem; }
.card { background: var(--card); border: 1px solid #eadfc8; border-radius: 14px; padding: 1rem; box-shadow: 0 8px 20px rgba(0,0,0,.06); }
h1 { margin: 0 0 .8rem 0; font-size: 1.4rem; }
.grid { display:grid; gap:.8rem; grid-template-columns: 1fr 1fr; }
.full { grid-column: 1 / -1; }
label { display:block; font-weight:600; margin-bottom:.3rem; }
input, textarea, button { width:100%; border-radius:10px; border:1px solid #d8ccb7; padding:.65rem; font-size:.95rem; }
textarea { min-height: 110px; resize: vertical; }
button { background: var(--accent); color:white; border:none; font-weight:700; cursor:pointer; }
button:disabled { opacity:.6; cursor:default; }
.muted { color: var(--muted); font-size:.88rem; }
.out img { width:100%; border-radius:12px; border:1px solid #e4d8c3; background:white; }
@media (max-width: 760px) { .grid { grid-template-columns: 1fr; } }
</style>
</head>
<body>
<div class=\"wrap\">
<div class=\"card\">
<h1>FLUX.2 Klein Playground</h1>
<p class=\"muted\">Private single-user UI at <code>/</code>. Distilled defaults: 4 steps, guidance 1.0.</p>
<form id=\"form\" class=\"grid\">
<div class=\"full\">
<label for=\"prompt\">Prompt</label>
<textarea id=\"prompt\" required>A cinematic photo portrait, natural skin texture, realistic lighting, high detail.</textarea>
</div>
<div>
<label for=\"width\">Width (multiple of 16)</label>
<input id=\"width\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
</div>
<div>
<label for=\"height\">Height (multiple of 16)</label>
<input id=\"height\" type=\"number\" value=\"1024\" step=\"16\" min=\"16\" />
</div>
<div>
<label for=\"steps\">Inference steps</label>
<input id=\"steps\" type=\"number\" value=\"4\" min=\"1\" max=\"100\" />
</div>
<div>
<label for=\"guidance\">Guidance scale</label>
<input id=\"guidance\" type=\"number\" value=\"1.0\" step=\"0.1\" />
</div>
<div>
<label for=\"seed\">Seed</label>
<input id=\"seed\" type=\"number\" value=\"0\" min=\"0\" />
</div>
<div>
<label for=\"edit\">Edit image (optional)</label>
<input id=\"edit\" type=\"file\" accept=\"image/*\" />
</div>
<div class=\"full\">
<button id=\"submit\" type=\"submit\">Generate</button>
<p id=\"status\" class=\"muted\">Idle</p>
</div>
</form>
</div>
<div class=\"card out\" style=\"margin-top:1rem\">
<h1>Output</h1>
<img id=\"preview\" alt=\"Generated image preview\" />
<p class=\"muted\">Tip: right-click image to save.</p>
</div>
</div>
<script>
const form = document.getElementById('form');
const statusEl = document.getElementById('status');
const preview = document.getElementById('preview');
const button = document.getElementById('submit');
async function fileToDataUrl(file) {
return await new Promise((resolve, reject) => {
const reader = new FileReader();
reader.onload = () => resolve(reader.result);
reader.onerror = reject;
reader.readAsDataURL(file);
});
}
form.addEventListener('submit', async (event) => {
event.preventDefault();
button.disabled = true;
statusEl.textContent = 'Generating...';
try {
const prompt = document.getElementById('prompt').value;
const width = Number(document.getElementById('width').value);
const height = Number(document.getElementById('height').value);
const steps = Number(document.getElementById('steps').value);
const guidance = Number(document.getElementById('guidance').value);
const seed = Number(document.getElementById('seed').value);
const file = document.getElementById('edit').files[0];
const imageDataUrl = file ? await fileToDataUrl(file) : null;
const payload = {
model: 'black-forest-labs/FLUX.2-klein-9B',
prompt,
size: `${width}x${height}`,
num_inference_steps: steps,
guidance_scale: guidance,
seed,
response_format: 'b64_json',
imageDataUrl,
};
const resp = await fetch('/v1/images/generations', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify(payload),
});
if (!resp.ok) {
const text = await resp.text();
throw new Error(text || `HTTP ${resp.status}`);
}
const data = await resp.json();
const b64 = data?.data?.[0]?.b64_json;
if (!b64) throw new Error('No image returned');
preview.src = `data:image/png;base64,${b64}`;
statusEl.textContent = `Done. Seed: ${data.seed ?? seed}`;
} catch (error) {
statusEl.textContent = `Error: ${error.message}`;
} finally {
button.disabled = false;
}
});
</script>
</body>
</html>
"""
@app.get("/health")
def health() -> dict[str, str]:
return {"status": "ok", "model": CHECKPOINT}
@app.post("/generate")
def generate(req: GenerateRequest) -> Response:
width, height = _normalize_dims(req.width, req.height)
seed = req.seed if not req.randomize_seed else random.randint(0, MAX_SEED)
image = _decode_image_data_url(req.image_data_url)
png_bytes, used_seed = _generate_png_bytes(
prompt=req.prompt,
width=width,
height=height,
num_inference_steps=req.num_inference_steps,
guidance_scale=req.guidance_scale,
seed=seed,
image=image,
)
return Response(
content=png_bytes,
media_type="image/png",
headers={"X-Seed": str(used_seed)},
)
@app.post("/v1/images/generations")
def compat_generate(req: CompatRequest) -> dict[str, Any]:
width, height = _parse_size(req.size, req.width, req.height)
seed_value = req.seed if req.seed is not None else 0
seed = seed_value if not req.randomize_seed else random.randint(0, MAX_SEED)
image = _decode_image_data_url(req.imageDataUrl)
png_bytes, used_seed = _generate_png_bytes(
prompt=req.prompt,
width=width,
height=height,
num_inference_steps=req.num_inference_steps or DEFAULT_STEPS,
guidance_scale=req.guidance_scale or DEFAULT_GUIDANCE,
seed=seed,
image=image,
)
b64_png = base64.b64encode(png_bytes).decode("ascii")
return {
"created": int(time.time()),
"model": CHECKPOINT,
"seed": used_seed,
"data": [{"b64_json": b64_png}],
}
if __name__ == "__main__":
import uvicorn
parser = argparse.ArgumentParser(description="Private FLUX.2 Klein FastAPI server")
parser.add_argument("--model", choices=["4b", "9b"], default=MODEL_SIZE)
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=6006)
args = parser.parse_args()
MODEL_SIZE = args.model
CHECKPOINT = _checkpoint_for(MODEL_SIZE)
uvicorn.run(app, host=args.host, port=args.port)
-68
View File
@@ -1,68 +0,0 @@
#!/usr/bin/env -S uv run
# /// script
# dependencies = ["httpx"]
# ///
import argparse
import concurrent.futures
import json
import re
import sys
import typing
from pathlib import Path
from typing import Optional
import logging
import httpx
def upload_image(image_path: Path) -> dict:
session = httpx.Client(timeout=20)
logging.info(f"uploading {image_path.name}")
with image_path.open("rb") as f:
res = session.post(
"https://fapping.empornium.sx/upload.php",
files={
"ImageUp": f,
},
)
res = session.get("https://fapping.empornium.sx/uploaded/", follow_redirects=True)
if m := re.search(r"var ImagesUp = (\[.+\]);", res.text, re.MULTILINE):
data: list[dict] = json.loads(m.group(1))
logging.info(f"upload result on the page={data}")
return dict(
image_url=data[0]["image_url"],
thumbnail_url=data[0]["image_thumb_url"],
)
logging.error(f"failed to find upload result on the page. html={res.text}")
raise ValueError(f"failed to parse upload result for {image_path.name}")
def parse_args(argv: Optional[typing.Sequence[str]] = None) -> argparse.Namespace:
arger = argparse.ArgumentParser()
arger.add_argument("image_paths", nargs="+", action="store", type=Path, help="Path to image to upload")
if not argv or len(argv) == 0:
arger.print_help()
sys.exit(1)
return arger.parse_args(argv)
def main():
args = parse_args(sys.argv[1:])
image_paths = args.image_paths
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
results = list(pool.map(upload_image, image_paths))
for it in results:
print(it["image_url"])
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
main()
-5
View File
@@ -1,5 +0,0 @@
module github.com/abdusco/playground
go 1.23.0
require golang.org/x/sync v0.10.0
-2
View File
@@ -1,2 +0,0 @@
golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ=
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
-133
View File
@@ -1,133 +0,0 @@
import http.server
import random
import socketserver
import sys
import threading
import time
import typing
import webbrowser
# language=html
inject_script = """
<script>
// Send periodic heartbeat
let heartbeatInterval = setInterval(() => {
fetch('/heartbeat', {method: 'POST'})
.catch(() => {}); // Ignore errors
}, 1000);
// Multiple event handlers for different close scenarios
function sendShutdown() {
clearInterval(heartbeatInterval);
// Synchronous request to ensure it gets sent
const xhr = new XMLHttpRequest();
xhr.open('POST', '/shutdown', false); // false makes it synchronous
try {
xhr.send();
} catch (e) {}
}
window.addEventListener('beforeunload', sendShutdown);
window.addEventListener('unload', sendShutdown);
// setTimeout(sendShutdown, 5000);
// Handle visibility change (tab hidden/visible)
document.addEventListener('visibilitychange', () => {
if (document.visibilityState === 'hidden') {
sendShutdown();
}
});
</script>
"""
def handler_for_html(html: str) -> typing.Type[http.server.SimpleHTTPRequestHandler]:
class Handler(http.server.SimpleHTTPRequestHandler):
def do_GET(self):
if self.path == "/":
self.send_response(200)
self.send_header("Content-type", "text/html")
self.end_headers()
self.wfile.write(html.encode())
else:
self.send_error(404)
def do_POST(self):
if self.path == "/shutdown":
self.send_response(200)
self.end_headers()
self.server.last_heartbeat = 0 # Force shutdown
self.server.shutdown()
elif self.path == "/heartbeat":
self.send_response(200)
self.end_headers()
self.server.last_heartbeat = time.time()
else:
self.send_error(404)
def log_message(self, format, *args):
# Suppress logging
pass
return Handler
class ServerWithHeartbeat(socketserver.TCPServer):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.last_heartbeat = time.time()
def monitor_heartbeat(server: ServerWithHeartbeat, timeout: int = 2):
while True:
time.sleep(0.5)
if time.time() - server.last_heartbeat > timeout:
print("No heartbeat received, shutting down...")
server.shutdown()
break
def run_server(handler: http.server.SimpleHTTPRequestHandler, port: int = 8000):
with ServerWithHeartbeat(("", port), handler) as httpd:
print(f"Serving at port {port}")
# Start heartbeat monitor in separate thread
monitor_thread = threading.Thread(target=monitor_heartbeat, args=(httpd,))
monitor_thread.daemon = True
monitor_thread.start()
httpd.serve_forever()
def main():
html = sys.stdin.read()
if "<body>" in html:
html_to_display = html.replace("</body>", f"{inject_script}</body>")
else:
html_to_display = f"<html><body>{html}{inject_script}</body></html>"
handler = handler_for_html(html_to_display)
# Start the server in a separate thread
port = random.randint(20000, 65535)
server_thread = threading.Thread(target=run_server, kwargs=dict(handler=handler, port=port))
server_thread.daemon = True
server_thread.start()
# Wait a moment for the server to start
time.sleep(0.5)
# Open the web browser
webbrowser.open(f"http://localhost:{port}")
# Wait for the server thread to finish
server_thread.join()
print("Server stopped")
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
print("\nReceived keyboard interrupt, exiting...")
raise SystemExit(0)
-119
View File
@@ -1,119 +0,0 @@
package main
import (
"context"
"fmt"
"io"
"log"
"net"
"net/http"
"os"
"os/exec"
"os/signal"
"runtime"
"strings"
"syscall"
"time"
)
// language=html
const injectScript = `
<script>
const sendShutdown = () => navigator.sendBeacon('/shutdown');
window.addEventListener('beforeunload', sendShutdown);
window.addEventListener('unload', sendShutdown);
</script>
`
func main() {
if err := run(); err != nil {
log.Fatal(err)
}
}
func newHandler(html string, onShutdown func()) *http.ServeMux {
handler := http.NewServeMux()
handler.HandleFunc("POST /shutdown", func(w http.ResponseWriter, r *http.Request) {
onShutdown()
})
handler.HandleFunc("GET /", func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/html")
w.Write([]byte(html))
})
return handler
}
func run() error {
html, err := io.ReadAll(os.Stdin)
if err != nil {
return fmt.Errorf("failed to read HTML from stdin: %w", err)
}
htmlStr := string(html)
if strings.Contains(htmlStr, "</body>") {
htmlStr = strings.ReplaceAll(string(html), "</body>", injectScript+"</body>")
} else {
htmlStr += injectScript
}
listener, err := net.Listen("tcp", ":0")
if err != nil {
return fmt.Errorf("failed to create listener: %w", err)
}
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer cancel()
port := listener.Addr().(*net.TCPAddr).Port
server := &http.Server{Addr: listener.Addr().String()}
server.Handler = newHandler(htmlStr, func() {
cancel()
})
servedURL := fmt.Sprintf("http://localhost:%d", port)
go func() {
log.Printf("serving at %s", servedURL)
if err := server.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Fatalf("failed to serve: %v", err)
}
}()
go func() {
// wait for server to start
time.Sleep(250 * time.Millisecond)
if err := openBrowser(servedURL); err != nil {
log.Fatalf("failed to open browser: %v", err)
}
}()
<-ctx.Done()
log.Printf("shutting down server")
if err := server.Shutdown(ctx); err != nil {
return fmt.Errorf("failed to shutdown server: %w", err)
}
return nil
}
func openBrowser(url string) error {
var cmd *exec.Cmd
switch runtime.GOOS {
case "darwin":
cmd = exec.Command("open", url)
case "linux":
cmd = exec.Command("xdg-open", url)
case "windows":
cmd = exec.Command("cmd", "/c", "start", url)
default:
return fmt.Errorf("invalid URL: %s", url)
}
return cmd.Run()
}
-220
View File
@@ -1,220 +0,0 @@
#!/usr/bin/env python3.9
import argparse
import dataclasses
import io
import itertools
import logging
import math
import statistics
import sys
import typing
from pathlib import Path
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
T = typing.TypeVar("T")
logger = logging.getLogger("imgsheet")
def chunk(it: typing.Iterable[T], size: int) -> typing.Iterable[typing.Iterable[T]]:
it = iter(it)
sentinel = ()
return iter(lambda: tuple(itertools.islice(it, size)), sentinel)
@dataclasses.dataclass
class PositionedImage:
path: Path
image_size: typing.Tuple[int, int]
size: typing.Tuple[int, int] = dataclasses.field(default_factory=lambda: (0, 0))
position: typing.Tuple[int, int] = dataclasses.field(default_factory=lambda: (0, 0))
@property
def top(self) -> int:
return self.position[1]
@property
def left(self) -> int:
return self.position[0]
@property
def bottom(self) -> int:
return self.height + self.top
@property
def right(self) -> int:
return self.width + self.left
@property
def width(self) -> int:
return self.size[0]
@property
def height(self) -> int:
return self.size[1]
def __post_init__(self):
self.size = self.image_size
def scale(self, scale: float) -> None:
self.size = (math.ceil(self.width * scale), math.ceil(self.height * scale))
class ImageSheet:
def __init__(
self,
width: int,
columns: int = 6,
font_path: typing.Optional[Path] = None,
label_format: typing.Optional[str] = None,
label_font_size: typing.Optional[int] = 10,
):
self.width = width
self.padding = 20
self.columns = columns
self.font_path = font_path
self.label_format = label_format
self.label_font_size = label_font_size
def create(self, image_paths: typing.Iterable[Path]) -> io.BytesIO:
images = []
for it in image_paths:
with Image.open(it) as img:
images.append(PositionedImage(it, img.size))
if self.font_path and self.font_path.is_file():
font = ImageFont.truetype(str(self.font_path), self.label_font_size)
else:
font = None
positioned = self._calculate_positions(images)
image_sheet = Image.new("RGB", (self.width, 0))
for row_images in positioned:
height = row_images[0].top + row_images[0].size[1]
extended_sheet = Image.new("RGB", (self.width, height + self.padding))
extended_sheet.paste(image_sheet, (0, 0))
image_sheet = extended_sheet
draw = ImageDraw.Draw(image_sheet)
for it in row_images:
img = Image.open(it.path)
img.thumbnail(it.size, Image.LANCZOS)
image_sheet.paste(img, it.position)
if font:
# write filename on image
text_x, text_y = it.left + 2, it.bottom + 2
filesize = it.path.stat().st_size / 1_048_576
text = self._render_text(
width=it.image_size[0],
height=it.image_size[1],
size_mb=filesize,
filename=it.path.name,
)
if text:
draw.text((text_x, text_y), text, (255, 255, 255), font=font, align="left", anchor="la")
print(f"{it.path.name}\t{it.position}\t{it.size}")
f = io.BytesIO()
image_sheet.save(f, format="JPEG")
return f
def _render_text(self, width: int, height: int, size_mb: float, filename: str) -> typing.Optional[str]:
if not self.label_format:
return None
return self.label_format.format(width=width, height=height, size_mb=size_mb, filename=filename)
def _calculate_positions(self, images: list[PositionedImage]) -> list[list[PositionedImage]]:
cols = self.columns
width = self.width
padding = self.padding
positioned = []
rows = list(chunk(images, cols))
x, y = 0, 0
for row, row_images in enumerate(rows):
row_images = list(row_images)
usable_width = width - (len(row_images) - 1) * padding
max_height = max(image.height for image in row_images)
# check if there are enough images to fill the row
if len(row_images) < cols:
mean_height = statistics.median(prow[0].height for prow in positioned)
max_height = mean_height
for img in row_images:
y_scaling = max_height / img.height
img.scale(y_scaling)
if len(row_images) >= cols:
total_width = sum(image.width for image in row_images)
x_scaling = usable_width / total_width
max_height = math.ceil(max_height * x_scaling)
for img in row_images:
img.scale(x_scaling)
is_last_row = row == len(rows) - 1
x = 0
for i, img in enumerate(row_images):
img.position = (x, y)
is_last_col = i == len(row_images) - 1
x = x + img.width + (0 if is_last_col else padding)
y = y + max_height + (0 if is_last_row else padding)
positioned.append(row_images)
return positioned
def parse_args(argv: typing.List[str]) -> argparse.Namespace:
arger = argparse.ArgumentParser(
description="Create a sheet of images", formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
arger.add_argument("--width", "-w", type=int, default=2000, help="Width of the sheet")
arger.add_argument("--columns", "-c", type=int, default=4, help="Number of columns")
arger.add_argument("--output", "-o", type=Path, default=Path("image_sheet.jpg"), help="Output file")
arger.add_argument("--font", dest="font_path", type=Path, help="Label font to use")
arger.add_argument("--font-size", dest="font_size", default=10, type=int, help="Label font size")
arger.add_argument(
"--label",
dest="label_format",
default="{width}x{height} {size_mb:.1f}MB {filename}",
help="Label format. Available tokens: width, height, size_mb, filename",
)
arger.add_argument("images", nargs="+", type=Path, help="Images or image directory to include in the sheet")
if len(argv) == 0:
arger.print_help()
sys.exit(1)
return arger.parse_args(argv)
def main():
args = parse_args(sys.argv[1:])
images: list[Path] = args.images
if len(images) == 1 and images[0].expanduser().is_dir():
images = sorted(it for it in images[0].resolve().glob("*.jpg"))
sheet = ImageSheet(
width=args.width,
columns=args.columns,
font_path=args.font_path,
label_font_size=args.font_size,
label_format=args.label_format,
)
sheet_img = sheet.create(images)
with args.output.open("wb") as f:
f.write(sheet_img.getvalue())
if __name__ == "__main__":
main()
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.
-35
View File
@@ -1,35 +0,0 @@
import json
import logging as logger
import sys
from pathlib import Path
import file_renamer
logger.basicConfig(
format=f"%(asctime)s {logger.BASIC_FORMAT}",
level=logger.DEBUG,
handlers=[
logger.FileHandler(Path(__file__).with_name("jd2_hook.log"), mode="a"),
logger.StreamHandler(),
],
)
def main():
try:
payload = json.loads(sys.argv[1])
except (IndexError, ValueError):
logger.error("Expected JSON as first argument")
return 1
for link in payload["links"]:
path = link["download_path"]
f = Path(path)
logger.info("")
if f.is_file():
file_renamer.fix_dates([f])
if __name__ == "__main__":
logger.debug(f"Got called with: {sys.argv}")
exit(main())
-98
View File
@@ -1,98 +0,0 @@
#!/usr/bin/env python3.9
import argparse
import concurrent.futures
import dataclasses
import re
import sys
import time
import typing
from pathlib import Path
from typing import Optional
import logging
from functools import cache
import httpx
@dataclasses.dataclass
class UploadResult:
image_url: str
thumbnail_url: str
session = httpx.Client(timeout=20)
@cache
def get_auth_token():
logging.info('extracting auth token')
res = session.get("https://jerking.empornium.ph/?agree-consent", follow_redirects=True)
try:
auth_token = re.search(r'name="auth_token" value="([^"]+)"', res.text).group(1)
except AttributeError:
auth_token = re.search(r'auth_token = "([^"]+)"', res.text).group(1)
logging.info(f'got {auth_token=}')
return auth_token
def upload_image(image_path: Path) -> UploadResult:
auth_token = get_auth_token()
logging.info(f'uploading {image_path.name}')
with image_path.open("rb") as f:
res = session.post(
"https://jerking.empornium.ph/json",
headers={"accept": "application/json"},
data={
"thumb_width": "400",
"thumb_height": "400",
"thumb_crop": "false",
"medium_width": "600",
"medium_crop": "false",
"type": "file",
"action": "upload",
"timestamp": str(round(time.time() * 1000)),
"auth_token": auth_token,
"nsfw": "1",
},
files={
"source": f,
},
)
data = res.json()
logging.info(f'got response={data}')
image_url = data["image"]["url"]
thumbnail_url = data["image"]["display_url"]
logging.info(f'uploaded {image_path.name}. url={image_url}')
return UploadResult(image_url, thumbnail_url)
def parse_args(argv: Optional[typing.Sequence[str]] = None) -> argparse.Namespace:
arger = argparse.ArgumentParser()
arger.add_argument('image_paths', nargs='+', action='store', type=Path, help='Path to image to upload')
if not argv or len(argv) == 0:
arger.print_help()
sys.exit(1)
return arger.parse_args(argv)
def main():
args = parse_args(sys.argv[1:])
image_paths = args.image_paths
get_auth_token()
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
results = list(pool.map(upload_image, image_paths))
for it in results:
print(it.image_url)
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
main()
-84
View File
@@ -1,84 +0,0 @@
#!/usr/bin/env python3
import argparse
import logging
import plistlib
import sys
import typing
from pathlib import Path
def parse_args(argv: list[str]) -> typing.Tuple[argparse.Namespace, list[str]]:
def parse_env(env: str) -> typing.Tuple[str, str]:
return env.split("=", maxsplit=1)
def parse_path(path: str) -> Path:
return Path(path).expanduser().resolve()
arger = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
arger.add_argument(
"-e", "--env", dest="env", metavar="ENVVAR", type=parse_env, nargs="+", help="Environment variables as K=V list"
)
arger.add_argument("-w", "--cwd", dest="cwd", default=Path.cwd(), type=parse_path, help="Working directory")
arger.add_argument("-l", "--label", required=True, help="Label")
arger.add_argument(
"--install",
action="store_true",
help="Install the plist to ~/Library/LaunchAgents/",
)
if not argv:
arger.print_help()
sys.exit(1)
args, unknown = arger.parse_known_args()
if unknown and unknown[0] == "--":
unknown = unknown[1:]
if not unknown:
arger.print_help()
sys.exit(1)
return args, unknown
def generate_plist(label: str, args: list[str], env: dict, cwd: Path) -> str:
doc = {
"Label": label,
"ProgramArguments": args,
"EnvironmentVariables": env,
"WorkingDirectory": str(cwd),
"RunAtLoad": True,
"StandardOutPath": str(cwd / f"{label}.stdout.log"),
"StandardErrorPath": str(cwd / f"{label}.stderr.log"),
}
return plistlib.dumps(
doc,
fmt=plistlib.FMT_XML,
sort_keys=True,
).decode()
def main():
args, unknown = parse_args(sys.argv[1:])
print(args, unknown)
plist_doc = generate_plist(
label=args.label,
args=unknown,
env=dict(args.env),
cwd=args.cwd,
)
print(plist_doc)
if args.install:
install_dir: Path = Path("~/Library/LaunchAgents/").expanduser()
logging.info(f"Installing plist to {install_dir}")
install_path = install_dir / f"{args.label}.plist"
install_path.write_text(plist_doc)
logging.info(f'Run "launchctl load -w {install_path}" to start the service')
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
main()
-336
View File
@@ -1,336 +0,0 @@
import itertools
import json
import logging
import pprint
import re
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import quote_plus
import boto3
import httpx
import urllib.parse
import subprocess
import typing
def list_repo_tags(*, repository_name: str, region: str, username: str, password: str) -> typing.List[dict]:
"""
Fetches the list of tags in a Codecommit repo using git ls-remote command
:returns: List of {tag, commit} items
"""
url = "https://{username}:{password}@git-codecommit.{region}.amazonaws.com/v1/repos/{repository_name}".format(
username=quote_plus(username),
password=quote_plus(password), # passwords are base64-encoded and might contain "/"
region=region,
repository_name=repository_name,
)
result = subprocess.run(["git", "ls-remote", "--refs", "--tags", url], text=True, stdout=subprocess.PIPE)
tags = []
for line in result.stdout.splitlines(keepends=False):
# lines are formatted as:
# 070161940636bab5d934af2702f4248e5a13eb29 refs/tags/ui_sandbox_acc104
# and columns are delimited with "\t"
if "refs/tags/" in line:
commit_hash, tag_ref = line.split("\t", maxsplit=1)
tag = tag_ref.replace("refs/tags/", "")
tags.append({"commit": commit_hash, "tag": tag})
# natural-sort by tag name
# flora_prod_1 < ... < flora_prod_9 > flora_prod_10
# vs lexical sort would give:
# flora_prod_1 < flora_prod_10 < ... < flora_prod_9
tags = sorted(tags, key=lambda t: [int(it) if it.isdigit() else it for it in re.split(r"(\d+)", t["tag"])])
return tags
def chunk(it, size):
it = iter(it)
sentinel = ()
return iter(lambda: tuple(itertools.islice(it, size)), sentinel)
def summarize_tags(tags: list[dict], take: int = 10) -> list[dict]:
re_prefix = re.compile(r"^(\D*)")
re_digits = re.compile(r"(\d+)")
grouped = defaultdict(list)
for t in tags:
tag_name = t["tag"]
if m := re_prefix.search(tag_name):
prefix = m.group(1)
grouped[prefix].append(t)
else:
print("no match", m)
flattened = [*grouped.pop("", [])]
for g, items in grouped.items():
sorted_items = sorted(
items, key=lambda t: [int(it) if it.isdigit() else it for it in re_digits.split(t["tag"])], reverse=True
)
flattened.extend(sorted_items[:take])
return flattened
abdus_creds = {
"username": "coder+1-at-400344683105",
"password": "4xdnGOtghKiT0M5CjulWf5raWwBSqT0Df2YkIY9Ws68=",
}
def get_commit_content():
client = boto3.client("codecommit")
res = client.get_file(
repositoryName="testing",
filePath="content.txt1",
commitSpecifier="vnext",
)
print(res)
def main():
import logging
from botocore.endpoint import Endpoint
logging.basicConfig(level=logging.INFO)
raw_make_request = Endpoint.make_request
def intercepted_make_request(*args, **kwargs):
aws_response, parsed_response = raw_make_request(*args, **kwargs)
logging.info('AWS response: %s', parsed_response)
return aws_response, parsed_response
Endpoint.make_request = intercepted_make_request
session = boto3.Session(
aws_access_key_id="AKIAQKZTNEVLBHDG5Q3X",
aws_secret_access_key="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",
region_name="eu-central-1",
)
cc = session.client("rds")
_ = cc.describe_db_clusters()
if __name__ == "__main__":
main()
exit()
# get_commit_content()
# flora b1cd47a221b84b2884b3d24198085a9f-f0dc2e747ff241649715c20d48232d8e
# oms 36363cc2f43b4cf5b1965af9f8925a4d-6deff091ce2d4c2483d8834c532a6f36
# omnitron 36363cc2f43b4cf5b1965af9f8925a4d-97400227a37940539c55489e1354ffc2
all_tags = list_repo_tags(
repository_name="b1cd47a221b84b2884b3d24198085a9f-f0dc2e747ff241649715c20d48232d8e",
username="sandbox-acc-api-user-at-023193265494",
password="sYPWKdhSx4JliXNI4ffRb4qWYBHLxA54nczSJ6CmYCY=",
region="eu-central-1",
)
print(all_tags)
exit(0)
all_tags.append({"tag": "1asd", "commit": "asdf"})
summarized = summarize_tags(all_tags, take=5)
commits = {t["commit"] for t in summarized}
session = boto3.Session(
aws_access_key_id="AKIAQKZTNEVLBHDG5Q3X",
aws_secret_access_key="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",
region_name="eu-central-1",
)
client = session.client("codecommit")
def get_commit_infos(commits: list[str]):
res = client.batch_get_commits(
repositoryName="36363cc2f43b4cf5b1965af9f8925a4d-97400227a37940539c55489e1354ffc2",
commitIds=commits,
)
return res["commits"]
with ThreadPoolExecutor() as pool:
results = list(pool.map(get_commit_infos, chunk(commits, 100)))
clone_ssh_url = repo["repositoryMetadata"]["cloneUrlSsh"]
print(repo)
exit(0)
# exit()
from botocore.awsrequest import AWSRequest
from botocore.signers import RequestSigner
# AWS Version 4 signing example
# DynamoDB API (CreateTable)
# See: http://docs.aws.amazon.com/general/latest/gr/sigv4_signing.html
# This version makes a POST request and passes request parameters
# in the body (payload) of the request. Auth information is passed in
# an Authorization header.
import sys, os, base64, datetime, hashlib, hmac
import requests # pip install requests
# ************* REQUEST VALUES *************
method = "POST"
service = "codecommit"
host = "codecommit.eu-central-1.amazonaws.com"
region = "eu-central-1"
endpoint = "https://codecommit.eu-central-1.amazonaws.com/"
# POST requests use a content type header. For DynamoDB,
# the content is JSON.
content_type = "application/x-amz-json-1.1"
# DynamoDB requires an x-amz-target header that has this format:
# DynamoDB_<API version>.<operationName>
amz_target = "CodeCommit_20150413.GetReferences"
# Request parameters for CreateTable--passed in a JSON block.
request_body = {"repositoryName": "36363cc2f43b4cf5b1965af9f8925a4d-6deff091ce2d4c2483d8834c532a6f36"}
request_parameters = json.dumps(request_body)
# Key derivation functions. See:
# http://docs.aws.amazon.com/general/latest/gr/signature-v4-examples.html#signature-v4-examples-python
def sign(key, msg):
return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()
def getSignatureKey(key, date_stamp, regionName, serviceName):
kDate = sign(("AWS4" + key).encode("utf-8"), date_stamp)
kRegion = sign(kDate, regionName)
kService = sign(kRegion, serviceName)
kSigning = sign(kService, "aws4_request")
return kSigning
# Read AWS access key from env. variables or configuration file. Best practice is NOT
# to embed credentials in code.
access_key = os.environ.get("AWS_ACCESS_KEY_ID", "AKIAQKZTNEVLBHDG5Q3X")
secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "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")
if access_key is None or secret_key is None:
print("No access key is available.")
sys.exit()
# Create a date for headers and the credential string
t = datetime.datetime.utcnow()
amz_date = t.strftime("%Y%m%dT%H%M%SZ")
date_stamp = t.strftime("%Y%m%d") # Date w/o time, used in credential scope
# ************* TASK 1: CREATE A CANONICAL REQUEST *************
# http://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
# Step 1 is to define the verb (GET, POST, etc.)--already done.
# Step 2: Create canonical URI--the part of the URI from domain to query
# string (use '/' if no path)
canonical_uri = "/"
## Step 3: Create the canonical query string. In this example, request
# parameters are passed in the body of the request and the query string
# is blank.
canonical_querystring = ""
# Step 4: Create the canonical headers. Header names must be trimmed
# and lowercase, and sorted in code point order from low to high.
# Note that there is a trailing \n.
canonical_headers = (
"content-type:"
+ content_type
+ "\n"
+ "host:"
+ host
+ "\n"
+ "x-amz-date:"
+ amz_date
+ "\n"
+ "x-amz-target:"
+ amz_target
+ "\n"
)
# Step 5: Create the list of signed headers. This lists the headers
# in the canonical_headers list, delimited with ";" and in alpha order.
# Note: The request can include any headers; canonical_headers and
# signed_headers include those that you want to be included in the
# hash of the request. "Host" and "x-amz-date" are always required.
# For DynamoDB, content-type and x-amz-target are also required.
signed_headers = "content-type;host;x-amz-date;x-amz-target"
# Step 6: Create payload hash. In this example, the payload (body of
# the request) contains the request parameters.
payload_hash = hashlib.sha256(request_parameters.encode("utf-8")).hexdigest()
# Step 7: Combine elements to create canonical request
canonical_request = (
method
+ "\n"
+ canonical_uri
+ "\n"
+ canonical_querystring
+ "\n"
+ canonical_headers
+ "\n"
+ signed_headers
+ "\n"
+ payload_hash
)
# ************* TASK 2: CREATE THE STRING TO SIGN*************
# Match the algorithm to the hashing algorithm you use, either SHA-1 or
# SHA-256 (recommended)
algorithm = "AWS4-HMAC-SHA256"
credential_scope = date_stamp + "/" + region + "/" + service + "/" + "aws4_request"
string_to_sign = (
algorithm
+ "\n"
+ amz_date
+ "\n"
+ credential_scope
+ "\n"
+ hashlib.sha256(canonical_request.encode("utf-8")).hexdigest()
)
# ************* TASK 3: CALCULATE THE SIGNATURE *************
# Create the signing key using the function defined above.
signing_key = getSignatureKey(secret_key, date_stamp, region, service)
# Sign the string_to_sign using the signing_key
signature = hmac.new(signing_key, (string_to_sign).encode("utf-8"), hashlib.sha256).hexdigest()
# ************* TASK 4: ADD SIGNING INFORMATION TO THE REQUEST *************
# Put the signature information in a header named Authorization.
authorization_header = (
algorithm
+ " "
+ "Credential="
+ access_key
+ "/"
+ credential_scope
+ ", "
+ "SignedHeaders="
+ signed_headers
+ ", "
+ "Signature="
+ signature
)
# For DynamoDB, the request can include any headers, but MUST include "host", "x-amz-date",
# "x-amz-target", "content-type", and "Authorization". Except for the authorization
# header, the headers must be included in the canonical_headers and signed_headers values, as
# noted earlier. Order here is not significant.
# # Python note: The 'host' header is added automatically by the Python 'requests' library.
headers = {
"User-Agent": "aws-sdk-js/2.627.0 promise",
"Content-Type": content_type,
"X-Amz-Date": amz_date,
"X-Amz-Target": amz_target,
"Authorization": authorization_header,
}
# ************* SEND THE REQUEST *************
print("\nBEGIN REQUEST++++++++++++++++++++++++++++++++++++")
print("Request URL = " + endpoint)
r = httpx.post(endpoint, json=request_body, headers=headers)
print("\nRESPONSE++++++++++++++++++++++++++++++++++++")
print("Response code: %d\n" % r.status_code)
print(r.text)
-438
View File
@@ -1,438 +0,0 @@
import dataclasses
import functools
import itertools
import logging
import re
import time
import typing
import uuid
from typing import Optional
import boto3
import botocore.exceptions
from botocore.exceptions import ClientError
logger = logging.getLogger(__name__)
def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
"""
Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
:param timeout: Timeout in seconds
:param backoff_seconds: Wait time between retries
:param bubble_errors: List of exceptions to bubble up
:return: Wrapped function
"""
bubble_errors = tuple(bubble_errors or [])
def wrapped(func: typing.Callable):
@functools.wraps(func)
def inner(*args, **kwargs):
waited = 0
while True:
try:
logger.debug(f"Calling {func}")
return func(*args, **kwargs)
except KeyboardInterrupt:
raise
except tuple(bubble_errors):
raise
except Exception as e:
doze = min(timeout - waited, backoff_seconds)
times = (timeout - waited) // backoff_seconds
logger.debug(f"{func} failed, will try {times} times in {doze} sec")
time.sleep(doze)
waited += backoff_seconds
if waited >= timeout:
raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
return inner
return wrapped
@dataclasses.dataclass
class RedisInstance:
cluster_id: str
host: str
port: int
arn: str
user_group_ids: typing.List[str] = dataclasses.field(default_factory=list)
def paginate_by_marker(
func: callable,
list_field: str,
) -> typing.Iterable:
"""
Simplifies paginating over AWS results with a simpler interface than Paginators.
:param func: A function that accepts a parameter named `Marker`
:param list_field: Result field to iterate on
"""
marker = ""
while True:
result = func(Marker=marker)
items = result.get(list_field, [])
yield from items
if "Marker" not in result:
break
marker = result["Marker"]
@dataclasses.dataclass
class RedisUser:
username: str
password: str
redis_key_prefix: str
class CreateUserResult(typing.TypedDict):
user_group_id: str
users: typing.List[RedisUser]
status: str
class RedisService:
def __init__(self, region: str):
self.elasticache = boto3.client("elasticache", region_name=region)
def create_redis_acl_user(
self,
replication_group_id: str,
users: typing.List[RedisUser],
) -> None:
"""
Creates a Redis ACL user. It will throw an error if the user group attached to the replication group is not ready.
This prevents us from quickly creating users. In that case, pass in a list of users instead.
It takes about 2 minutes until the user is ready.
"""
repl_groups = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)
user_group_id: str = repl_groups['ReplicationGroups'][0]['UserGroupIds'][0]
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
"""
Limits imposed by AWS:
- User groups per replication group = 1
- Users per user group = 100
- Number of users = 1000
- Number of user groups = 100
https://docs.aws.amazon.com/AmazonElastiCache/latest/red-ug/Clusters.RBAC.html#Users-groups-to-RGs
In short, we can't place more than 100 users on a single Redis instance.
"""
total_users = len(group['UserIds'])
if total_users >= 100:
raise Exception('Replication group reached the limit of 100 users')
self._update_or_create_user_group(user_group_id, users)
def delete_redis_acl_user(self, username: str) -> None:
logger.info(f'Deleting user {username}')
self.elasticache.delete_user(UserId=username)
# Even though we just deleted the user, it takes a couple of seconds until AWS updates user group status.
# So, we have to wait a bit before we can actually wait and check that the changes have propagated.
# Overall, it should take about 2-3 minutes to delete a user.
def create_redis_acl(
self,
*,
replication_group_id: str,
description: str = None,
node_type: str,
cache_subnet_group_name: str,
security_group_id: str,
tags: typing.Dict[str, str] = None,
users: typing.List[RedisUser],
) -> dict:
user_group_id = f"{replication_group_id}-ug"
# we have to wait (~1m) until the user group is ready before we can create the ACL
self._update_or_create_user_group(user_group_id, users, wait=True)
cache_param_group = 'redis-acl-with-100-db'
self._ensure_cache_parameter_group(cache_param_group)
logger.info(f"Creating Redis replication group {replication_group_id=}")
default_replication_group_kwargs = dict(
ReplicasPerNodeGroup=0,
Engine="redis",
EngineVersion="6.x",
TransitEncryptionEnabled=True, # must be True to use Redis ACL
MultiAZEnabled=False,
AutomaticFailoverEnabled=False,
)
repl_result = self.elasticache.create_replication_group(
**default_replication_group_kwargs,
ReplicationGroupId=replication_group_id,
ReplicationGroupDescription=description or f'{replication_group_id} replication group',
CacheNodeType=node_type,
CacheSubnetGroupName=cache_subnet_group_name,
CacheParameterGroupName=cache_param_group,
UserGroupIds=[
user_group_id,
],
SecurityGroupIds=[
security_group_id,
],
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
)
# replication group takes ~8m minutes to be ready
return {
"replication_group_id": repl_result["ReplicationGroup"]["ReplicationGroupId"],
'arn': repl_result["ReplicationGroup"]["ARN"],
'status': repl_result["ReplicationGroup"]["Status"],
'users': users,
'_response': repl_result,
}
def _update_or_create_user_group(
self, user_group_id: str, users: typing.List[RedisUser], wait: bool = False
) -> None:
"""
Creates or updates a user group with the given users. Waits until the user group is active.
:param user_group_id:
:param users:
"""
# we have to add the `default` user for backwards compatibility
existing_users = self.elasticache.describe_users(
Filters=[{"Name": "user-id", "Values": [u.username for u in users]}],
)["Users"]
existing_user_ids = [it["UserId"] for it in existing_users]
created_user_ids = ["default"]
for user in users:
if user.username in existing_user_ids:
logger.info(f"User {user.username} is already exists")
created_user_ids.append(user.username)
continue
redis_acl = f"on +@all -@dangerous ~{user.redis_key_prefix}*"
logger.info(f"Creating tenant {user.username=}")
user_result = self.elasticache.create_user(
UserId=user.username,
UserName=user.username,
Passwords=[user.password],
AccessString=redis_acl,
Engine="redis",
)
created_user_ids.append(user_result["UserId"])
try:
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
member_user_ids: typing.List[str] = group["UserIds"]
users_to_add = set(created_user_ids) - set(member_user_ids)
if users_to_add:
self.elasticache.modify_user_group(
UserGroupId=user_group_id,
UserIdsToAdd=list(users_to_add),
)
# this will take some time (~45s) until the changes propagate
except self.elasticache.exceptions.UserGroupNotFoundFault:
logger.info('User group does not exist, creating')
_ = self.elasticache.create_user_group(
UserGroupId=user_group_id,
Engine="redis",
UserIds=created_user_ids,
)
# user group creation takes around 60s
if wait:
self._wait_user_group(user_group_id)
@retry(timeout=60, backoff_seconds=5)
def _wait_user_group(self, user_group_id: str) -> None:
logger.info(f'Checking status of user group {user_group_id=}')
g = self.elasticache.describe_user_groups(UserGroupId=user_group_id)["UserGroups"][0]
assert g["Status"] == "active"
def is_redis_acl_ready(self, replication_group_id: str) -> bool:
repl = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
'ReplicationGroups'
][0]
is_redis_available = repl['Status'] == 'available'
user_group_id = repl['UserGroupIds'][0]
group = self.elasticache.describe_user_groups(UserGroupId=user_group_id)['UserGroups'][0]
is_group_available = group['Status'] == 'active'
return is_group_available and is_redis_available
def get_redis_acl(self, replication_group_id: str) -> dict:
redis = self.elasticache.describe_replication_groups(ReplicationGroupId=replication_group_id)[
'ReplicationGroups'
][0]
return dict(
replication_group_id=redis['ReplicationGroupId'],
arn=redis['ARN'],
host=redis["NodeGroups"][0]["PrimaryEndpoint"]["Address"],
port=redis["NodeGroups"][0]["PrimaryEndpoint"]["Port"],
)
def delete_redis_acl(self, replication_group_id: str) -> None:
self.elasticache.delete_replication_group(ReplicationGroupId=replication_group_id)
def _ensure_cache_parameter_group(self, parameter_group_name: str) -> None:
try:
_ = self.elasticache.describe_cache_parameter_groups(CacheParameterGroupName=parameter_group_name)[
'CacheParameterGroups'
][0]
return
except self.elasticache.exceptions.CacheParameterGroupNotFoundFault:
pass
logger.info(f"Creating cache parameter group {parameter_group_name=}")
_ = self.elasticache.create_cache_parameter_group(
CacheParameterGroupName=parameter_group_name,
CacheParameterGroupFamily="redis6.x",
Description="Redis parameter group for the Redis ACL",
)
_ = self.elasticache.modify_cache_parameter_group(
CacheParameterGroupName=parameter_group_name,
ParameterNameValues=[
{
"ParameterName": "databases",
"ParameterValue": "100",
},
],
)
def akinon_create_redis(
*,
k8s_cluster_name: str,
region: str,
node_type: str,
replicas: int,
owner_arn: str,
app_name: str,
role: str,
):
subnet_group_name = f"{k8s_cluster_name}-redis-subg"
security_group_name = f"{k8s_cluster_name}-redis-sg"
# owner arn is formatted like:
# arn:aws:iam::412344683105:user/myusername
username = owner_arn.split("/")[1]
tags = {
"CostCenter": f"{username}-{app_name}-{role}-redis",
}
cache_cluster_id = uuid.uuid4().hex
r = RedisService(region)
return r.create_redis(
cluster_id=cache_cluster_id,
node_type=node_type,
replicas=replicas,
security_group_name=security_group_name,
subnet_group_name=subnet_group_name,
tags=tags,
)
def akinon_create_redis_with_acl(
k8s_cluster_name: str,
owner_arn: str,
node_type: str,
app_name: str,
role: str,
users: typing.List[RedisUser],
) -> RedisInstance:
cache_cluster_id = uuid.uuid4().hex
subnet_group_name = f"{k8s_cluster_name}-redis-subg"
security_group_name = f"{k8s_cluster_name}-redis-sg"
# owner arn is formatted as:
# arn:aws:iam::412344683105:user/myusername
username = owner_arn.split("/")[1]
tags = {
"CostCenter": f"{username}-{app_name}-{role}-redis",
}
return r.create_redis_acl(
cluster_id=cache_cluster_id,
node_type=node_type,
subnet_group_name=subnet_group_name,
security_group_name=security_group_name,
tags=tags,
users=users,
)
def chunk(it, size):
it = iter(it)
sentinel = ()
return iter(lambda: tuple(itertools.islice(it, size)), sentinel)
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
r = RedisService(region="eu-central-1")
# res = r.create_redis_acl(
# replication_group_id="zerodev3",
# description="zero redis dev",
# node_type="cache.t3.micro",
# cache_subnet_group_name="dev",
# security_group_id="sg-c11a77ac",
# tags={"CostCenter": "zero123"},
# users=[RedisUser(username="u1", password="zeropassword1234", redis_key_prefix="zero:")],
# )
# print(res)
r.create_redis_acl_user(
replication_group_id='zerodev3',
users=[
RedisUser(
username='u2',
password='testtesttesttest123',
redis_key_prefix='test:',
)
],
)
# r.delete_redis_acl_user('test6')
while True:
if r.is_redis_acl_ready('zerodev3'):
logger.info('ready')
print(r.get_redis_acl('zerodev3'))
break
logger.info('not ready')
time.sleep(10)
exit()
res = r.create_redis_acl(
replication_group_id="zerodev",
description="zero redis dev",
node_type="cache.t3.micro",
cache_subnet_group_name="dev",
security_group_id="sg-c11a77ac",
tags={"CostCenter": "zero123"},
users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
)
print(res)
#
# redis = r.get_redis_acl(cluster_arn="arn:aws:elasticache:eu-central-1:400344683105:replicationgroup:zero")
# r.add_redis_user(
# cluster_id=redis.cluster_id,
# users=[RedisUser(username="zero27", password="zero27password123", redis_key_prefix="zero27:")],
# )
# print(r.add_redis_user(cluster_id="zero1", users=[RedisUser("zero26", "zero26password123", "zero26:")]))
#
# print(
# r.create_redis_with_acl(
# cluster_id="zero",
# cluster_description="zero redis",
# node_type="cache.t3.micro",
# subnet_group_name="testing-subnet",
# security_group_name="default",
# tags={"CostCenter": "zero123"},
# users=[RedisUser(username="zerouser", password="zeropassword1234", redis_key_prefix="zero:")],
# )
# )
-238
View File
@@ -1,238 +0,0 @@
import dataclasses
import datetime
import functools
import json
import logging
import time
import typing
import uuid
import boto3
import botocore
import requests
from botocore.client import BaseClient
from botocore.errorfactory import BaseClientExceptions
from botocore.exceptions import ClientError
logger = logging.getLogger(__name__)
# fmt: off
_response_describe_elasticsearch_domain_success = {'ResponseMetadata': {'RequestId': '38d7b555-2699-4868-9139-c47b1c2d70db', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amzn-requestid': '38d7b555-2699-4868-9139-c47b1c2d70db', 'content-type': 'application/json', 'content-length': '2314', 'date': 'Tue, 11 Jan 2022 07:52:24 GMT'}, 'RetryAttempts': 0}, 'DomainStatus': {'DomainId': '400344683105/es-d8a24850a0ae4b82af26', 'DomainName': 'es-d8a24850a0ae4b82af26', 'ARN': 'arn:aws:es:eu-central-1:400344683105:domain/es-d8a24850a0ae4b82af26', 'Created': True, 'Deleted': False, 'Endpoints': {'vpc': 'vpc-es-d8a24850a0ae4b82af26-gcw46l5kiiy6oylbhpqpr7s454.eu-central-1.es.amazonaws.com'}, 'Processing': False, 'UpgradeProcessing': False, 'ElasticsearchVersion': '7.8', 'ElasticsearchClusterConfig': {'InstanceType': 't2.medium.elasticsearch', 'InstanceCount': 2, 'DedicatedMasterEnabled': False, 'ZoneAwarenessEnabled': True, 'ZoneAwarenessConfig': {'AvailabilityZoneCount': 2}, 'WarmEnabled': False, 'ColdStorageOptions': {'Enabled': False}}, 'EBSOptions': {'EBSEnabled': True, 'VolumeType': 'gp2', 'VolumeSize': 10}, 'AccessPolicies': '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":"*","Action":"es:*","Resource":"*"}]}', 'SnapshotOptions': {'AutomatedSnapshotStartHour': 0}, 'VPCOptions': {'VPCId': 'vpc-8c28b5e7', 'SubnetIds': ['subnet-004909d95a03c5fcd', 'subnet-b8f596f5'], 'AvailabilityZones': ['eu-central-1a', 'eu-central-1c'], 'SecurityGroupIds': ['sg-c11a77ac']}, 'CognitoOptions': {'Enabled': False}, 'EncryptionAtRestOptions': {'Enabled': False}, 'NodeToNodeEncryptionOptions': {'Enabled': False}, 'AdvancedOptions': {'override_main_response_version': 'false', 'rest.action.multi.allow_explicit_index': 'true'}, 'ServiceSoftwareOptions': {'CurrentVersion': 'R20211203-P2', 'NewVersion': '', 'UpdateAvailable': False, 'Cancellable': False, 'UpdateStatus': 'COMPLETED', 'Description': 'There is no software update available for this domain.', 'AutomatedUpdateDate': datetime.datetime(2021, 12, 14, 3, 38, 38), 'OptionalDeployment': False}, 'DomainEndpointOptions': {'EnforceHTTPS': True, 'TLSSecurityPolicy': 'Policy-Min-TLS-1-0-2019-07', 'CustomEndpointEnabled': False}, 'AdvancedSecurityOptions': {'Enabled': False, 'InternalUserDatabaseEnabled': False}, 'AutoTuneOptions': {'State': 'ENABLE_IN_PROGRESS'}}}
_response_list_domain_names_success = {'ResponseMetadata': {'RequestId': '25923514-a426-4223-939a-4436f1d6e243', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amzn-requestid': '25923514-a426-4223-939a-4436f1d6e243', 'content-type': 'application/json', 'content-length': '87', 'date': 'Tue, 11 Jan 2022 07:53:37 GMT'}, 'RetryAttempts': 0}, 'DomainNames': [{'DomainName': 'es-d8a24850a0ae4b82af26', 'EngineType': 'Elasticsearch'}]}
_response_describe_elasticsearch_domains_success = {'ResponseMetadata': {'RequestId': '61cfe5f4-42e7-4233-bca9-53f0ca62099d', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amzn-requestid': '61cfe5f4-42e7-4233-bca9-53f0ca62099d', 'content-type': 'application/json', 'content-length': '2320', 'date': 'Tue, 11 Jan 2022 07:55:25 GMT'}, 'RetryAttempts': 0}, 'DomainStatusList': [{'DomainId': '400344683105/es-d8a24850a0ae4b82af26', 'DomainName': 'es-d8a24850a0ae4b82af26', 'ARN': 'arn:aws:es:eu-central-1:400344683105:domain/es-d8a24850a0ae4b82af26', 'Created': True, 'Deleted': False, 'Endpoints': {'vpc': 'vpc-es-d8a24850a0ae4b82af26-gcw46l5kiiy6oylbhpqpr7s454.eu-central-1.es.amazonaws.com'}, 'Processing': False, 'UpgradeProcessing': False, 'ElasticsearchVersion': '7.8', 'ElasticsearchClusterConfig': {'InstanceType': 't2.medium.elasticsearch', 'InstanceCount': 2, 'DedicatedMasterEnabled': False, 'ZoneAwarenessEnabled': True, 'ZoneAwarenessConfig': {'AvailabilityZoneCount': 2}, 'WarmEnabled': False, 'ColdStorageOptions': {'Enabled': False}}, 'EBSOptions': {'EBSEnabled': True, 'VolumeType': 'gp2', 'VolumeSize': 10}, 'AccessPolicies': '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":"*","Action":"es:*","Resource":"*"}]}', 'SnapshotOptions': {'AutomatedSnapshotStartHour': 0}, 'VPCOptions': {'VPCId': 'vpc-8c28b5e7', 'SubnetIds': ['subnet-004909d95a03c5fcd', 'subnet-b8f596f5'], 'AvailabilityZones': ['eu-central-1a', 'eu-central-1c'], 'SecurityGroupIds': ['sg-c11a77ac']}, 'CognitoOptions': {'Enabled': False}, 'EncryptionAtRestOptions': {'Enabled': False}, 'NodeToNodeEncryptionOptions': {'Enabled': False}, 'AdvancedOptions': {'override_main_response_version': 'false', 'rest.action.multi.allow_explicit_index': 'true'}, 'ServiceSoftwareOptions': {'CurrentVersion': 'R20211203-P2', 'NewVersion': '', 'UpdateAvailable': False, 'Cancellable': False, 'UpdateStatus': 'COMPLETED', 'Description': 'There is no software update available for this domain.', 'AutomatedUpdateDate': datetime.datetime(2021, 12, 14, 3, 38, 38), 'OptionalDeployment': False}, 'DomainEndpointOptions': {'EnforceHTTPS': True, 'TLSSecurityPolicy': 'Policy-Min-TLS-1-0-2019-07', 'CustomEndpointEnabled': False}, 'AdvancedSecurityOptions': {'Enabled': False, 'InternalUserDatabaseEnabled': False}, 'AutoTuneOptions': {'State': 'ENABLE_IN_PROGRESS'}}]}
_response_delete_elasticsearch_domain_not_found = ClientError({'Error': {'Message': 'Domain not found: asd', 'Code': 'ResourceNotFoundException'}, 'ResponseMetadata': {'RequestId': 'd0669560-77f6-424c-bbc6-fb5878daa494', 'HTTPStatusCode': 409, 'HTTPHeaders': {'x-amzn-requestid': 'd0669560-77f6-424c-bbc6-fb5878daa494', 'x-amzn-errortype': 'ResourceNotFoundException', 'content-type': 'application/json', 'content-length': '35', 'date': 'Tue, 11 Jan 2022 08:06:07 GMT'}, 'RetryAttempts': 0}}, 'DeleteElasticsearchDomain')
_response_delete_elasticsearch_domain_success = {'ResponseMetadata': {'RequestId': '1e9046f4-c4ee-4fda-8dfe-dd23f0161b7e', 'HTTPStatusCode': 200, 'HTTPHeaders': {'x-amzn-requestid': '1e9046f4-c4ee-4fda-8dfe-dd23f0161b7e', 'content-type': 'application/json', 'content-length': '2312', 'date': 'Tue, 11 Jan 2022 08:31:33 GMT'}, 'RetryAttempts': 0}, 'DomainStatus': {'DomainId': '400344683105/es-d8a24850a0ae4b82af26', 'DomainName': 'es-d8a24850a0ae4b82af26', 'ARN': 'arn:aws:es:eu-central-1:400344683105:domain/es-d8a24850a0ae4b82af26', 'Created': True, 'Deleted': True, 'Endpoints': {'vpc': 'vpc-es-d8a24850a0ae4b82af26-gcw46l5kiiy6oylbhpqpr7s454.eu-central-1.es.amazonaws.com'}, 'Processing': True, 'UpgradeProcessing': False, 'ElasticsearchVersion': '7.8', 'ElasticsearchClusterConfig': {'InstanceType': 't2.medium.elasticsearch', 'InstanceCount': 2, 'DedicatedMasterEnabled': False, 'ZoneAwarenessEnabled': True, 'ZoneAwarenessConfig': {'AvailabilityZoneCount': 2}, 'WarmEnabled': False, 'ColdStorageOptions': {'Enabled': False}}, 'EBSOptions': {'EBSEnabled': True, 'VolumeType': 'gp2', 'VolumeSize': 10}, 'AccessPolicies': '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":"*","Action":"es:*","Resource":"*"}]}', 'SnapshotOptions': {'AutomatedSnapshotStartHour': 0}, 'VPCOptions': {'VPCId': 'vpc-8c28b5e7', 'SubnetIds': ['subnet-004909d95a03c5fcd', 'subnet-b8f596f5'], 'AvailabilityZones': ['eu-central-1a', 'eu-central-1c'], 'SecurityGroupIds': ['sg-c11a77ac']}, 'CognitoOptions': {'Enabled': False}, 'EncryptionAtRestOptions': {'Enabled': False}, 'NodeToNodeEncryptionOptions': {'Enabled': False}, 'AdvancedOptions': {'override_main_response_version': 'false', 'rest.action.multi.allow_explicit_index': 'true'}, 'ServiceSoftwareOptions': {'CurrentVersion': 'R20211203-P2', 'NewVersion': '', 'UpdateAvailable': False, 'Cancellable': False, 'UpdateStatus': 'COMPLETED', 'Description': 'There is no software update available for this domain.', 'AutomatedUpdateDate': datetime.datetime(2021, 12, 14, 3, 38, 38), 'OptionalDeployment': False}, 'DomainEndpointOptions': {'EnforceHTTPS': True, 'TLSSecurityPolicy': 'Policy-Min-TLS-1-0-2019-07', 'CustomEndpointEnabled': False}, 'AdvancedSecurityOptions': {'Enabled': False, 'InternalUserDatabaseEnabled': False}, 'AutoTuneOptions': {'State': 'ENABLE_IN_PROGRESS'}}}
# fmt: on
def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
"""
Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
:param timeout: Timeout in seconds
:param backoff_seconds: Wait time between retries
:param bubble_errors: List of exceptions to bubble up
:return: Wrapped function
"""
bubble_errors = tuple(bubble_errors or [])
def wrapped(func: typing.Callable):
@functools.wraps(func)
def func_retrier(*args, **kwargs):
waited = 0
started_at = datetime.datetime.now()
while True:
try:
logger.debug(f"Calling {func}")
result = func(*args, **kwargs)
logger.debug(f"{func} completed in {(datetime.datetime.now() - started_at)}")
return result
except KeyboardInterrupt:
raise
except tuple(bubble_errors):
raise
except Exception as e:
doze = min(timeout - waited, backoff_seconds)
times = (timeout - waited) // backoff_seconds
logger.debug(f"{func} failed, will try {times} times in {doze} sec", exc_info=True)
time.sleep(doze)
waited += backoff_seconds
if waited >= timeout:
raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
return func_retrier
return wrapped
@dataclasses.dataclass
class ElasticsearchInstance:
domain_name: str
endpoint: str
arn: str
class ElasticsearchService:
domain_name_format = "es-{}"
default_version = "7.8"
valid_versions = {"7.10", "7.9", "7.8", "7.7", "7.4", "7.1", "6.8", "6.7", "6.5", "6.4", "6.3", "6.2", "6.0", "5.6", "5.5"} # fmt: skip
def __init__(self, region: str):
self.es = boto3.client("es", region_name=region)
def create_instance(
self,
node_type: str,
total_nodes: int,
subnet_group_ids: typing.List[str],
security_group_id: str,
version: str = None,
tags: dict = None,
volume_size_gb: int = 30,
total_availability_zones: int = 3,
) -> ElasticsearchInstance:
if not version:
version = self.default_version
assert version in self.valid_versions, "Invalid version"
domain_name = self.domain_name_format.format(uuid.uuid4().hex[:20])
logger.info(
"Creating a new Elasticsearch domain",
extra=dict(version=version, domain_name=domain_name, node_type=node_type, total_nodes=total_nodes),
)
es_result = self.es.create_elasticsearch_domain(
DomainName=domain_name,
ElasticsearchVersion=version,
ElasticsearchClusterConfig={
"InstanceType": node_type,
"InstanceCount": total_nodes,
"ZoneAwarenessEnabled": True,
"ZoneAwarenessConfig": {
"AvailabilityZoneCount": total_availability_zones,
},
},
VPCOptions={
"SubnetIds": subnet_group_ids,
"SecurityGroupIds": [security_group_id],
},
EBSOptions={
"EBSEnabled": True,
"VolumeType": "gp2",
"VolumeSize": volume_size_gb,
},
AccessPolicies=json.dumps(
{
"Version": "2012-10-17",
"Statement": [{"Effect": "Allow", "Principal": "*", "Action": "es:*", "Resource": "*"}],
}
),
DomainEndpointOptions={"EnforceHTTPS": True},
TagList=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
)
logger.info("Waiting until ES domain becomes available. This will take a while (~15 min)")
return self._find_instance(domain_name)
@retry(timeout=20 * 60, backoff_seconds=30)
def _find_instance(self, domain_name: str) -> ElasticsearchInstance:
result = self.es.describe_elasticsearch_domain(DomainName=domain_name)["DomainStatus"]
return ElasticsearchInstance(endpoint=result["Endpoints"]["vpc"], arn=result["ARN"], domain_name=domain_name)
def find_domain_name_by_arn(self, arn: str):
res = self.es.list_domain_names(EngineType='Elasticsearch')
domain_names = [it['DomainName'] for it in res['DomainNames']]
res = self.es.describe_elasticsearch_domains(DomainNames=domain_names)
domain_name_by_arn = {it['ARN']: it['DomainName'] for it in res['DomainStatusList']}
print(res)
def main():
username, password = 'bgNNs4EWPTsuQJ24', '3asinFjikT5MgwRz!'
"""
GET _search
{
"query": {
"match_all": {}
}
}
###
PUT /my-index
PUT /my-index2
PUT /my-index3
PUT /my-index3-4
###
DELETE /my-index*?expand_wildcards=all
###
DELETE /my-index1
###
GET /my-index
GET /_stats/indexing
###
GET /_cat/indices?v
###
GET /kibana_sample_data_ecommerce/
###
GET /_all/_mapping
## delete multiple
PUT /my-index
PUT /my-index2
DELETE /my-index,my-index2
GET /my-index
###
GET /_aliases
###
GET /_stats/
user: qdSmWKS8p3VrUyBC
pwd: un6f2vTrRjoa8orn!
"""
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
logging.getLogger("botocore").setLevel(logging.INFO)
# ElasticsearchService('eu-central-1')
e = ElasticsearchService("eu-central-1")
esi = ElasticsearchInstance(
endpoint='vpc-es-d8a24850a0ae4b82af26-gcw46l5kiiy6oylbhpqpr7s454.eu-central-1.es.amazonaws.com',
arn='arn:aws:es:eu-central-1:400344683105:domain/es-d8a24850a0ae4b82af26',
domain_name='es-d8a24850a0ae4b82af26',
)
# e.find_instance_by_arn(esi.arn)
# exit()
try:
es = boto3.client('es')
res = es.delete_elasticsearch_domain(DomainName=esi.domain_name)
print(res)
except ClientError as e:
error_code = e.__class__.__name__
print(error_code)
e = ElasticsearchService("eu-central-1")
# print(
# e.create_instance(
# node_type="t2.medium.elasticsearch",
# total_nodes=2,
# volume_size_gb=10,
# total_availability_zones=2,
# security_group_id="sg-c11a77ac",
# subnet_group_ids=["subnet-b8f596f5", "subnet-004909d95a03c5fcd"],
# tags={"CostCenter": "my"},
# )
# )
-401
View File
@@ -1,401 +0,0 @@
import dataclasses
import datetime
import functools
import logging
import time
import typing
import uuid
import boto3
# import mysql.connector.connection
import psycopg2
import pymysql
from psycopg2 import sql
import psycopg2.extensions
logger = logging.getLogger(__name__)
@dataclasses.dataclass
class Credentials:
username: str
password: str
@classmethod
def new(cls):
username_format = "u{}" # must start with a letter
return cls(
username=username_format.format(uuid.uuid4().hex),
password=uuid.uuid4().hex,
) # make sure it starts with a letter
@dataclasses.dataclass
class DbConnection:
endpoint: str
port: int
@dataclasses.dataclass
class DbInstance(DbConnection):
cluster_id: str
arn: str
reader_endpoint: typing.Optional[str] = None
database_name: typing.Optional[str] = None
user: typing.Optional[Credentials] = None
master_user: typing.Optional[Credentials] = None
def retry(timeout: int, backoff_seconds: int = 10, bubble_errors: typing.List[typing.Type[Exception]] = None):
"""
Tries calling a function {timeout} seconds until it succeeds or gives up and throw an error
:param timeout: Timeout in seconds
:param backoff_seconds: Wait time between retries
:param bubble_errors: List of exceptions to bubble up
:return: Wrapped function
"""
bubble_errors = tuple(bubble_errors or [])
def wrapped(func: typing.Callable):
@functools.wraps(func)
def func_retrier(*args, **kwargs):
waited = 0
started_at = datetime.datetime.now()
while True:
try:
logger.debug(f"Calling {func}")
result = func(*args, **kwargs)
logger.debug(f"{func} completed in {(datetime.datetime.now() - started_at)}")
return result
except KeyboardInterrupt:
raise
except tuple(bubble_errors):
raise
except Exception as e:
doze = min(timeout - waited, backoff_seconds)
times = (timeout - waited) // backoff_seconds
logger.debug(f"{func} failed, will try {times} times in {doze} sec", exc_info=True)
time.sleep(doze)
waited += backoff_seconds
if waited >= timeout:
raise TimeoutError(f"Couldn't get a successful result from {func} in {timeout} seconds") from e
return func_retrier
return wrapped
def paginate_by_marker(
func: callable,
list_field: str,
) -> typing.Iterable:
"""
Simplifies paginating over AWS results with a simpler interface than Paginators.
:param func: A function that accepts a parameter named `Marker`
:param list_field: Result field to iterate on
"""
marker = ""
while True:
result = func(Marker=marker)
items = result.get(list_field, [])
yield from items
if "Marker" not in result:
break
marker = result["Marker"]
class PostgresqlService:
db_type = "Postgresql"
db_engine = "aurora-postgresql"
db_engine_version = "10.11"
default_master_database_name = "postgres"
cluster_id_format = "pg-{}" # must start with a letter
tenant_database_name_format = "db{}" # must start with a letter
def __init__(self, region: str):
self.rds = boto3.client("rds", region_name=region)
def create_db(
self,
total_instances: int,
db_subnet_group_name: str,
node_type: str,
security_group_id: str,
engine_version: str = None,
backup_retention_days: int = 30,
public: bool = False,
tags: dict = None,
):
if not engine_version:
engine_version = self.db_engine_version
cluster_id = self.cluster_id_format.format(uuid.uuid4().hex)
master_creds = Credentials.new()
logger.info(
f"Creating {self.db_type} cluster",
extra=dict(
cluster_id=cluster_id,
engine_version=engine_version,
db_subnet_group_name=db_subnet_group_name,
security_group_id=security_group_id,
),
)
db_result = self.rds.create_db_cluster(
DBClusterIdentifier=cluster_id,
Engine=self.db_engine,
EngineVersion=engine_version,
MasterUsername=master_creds.username,
MasterUserPassword=master_creds.password,
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
DBSubnetGroupName=db_subnet_group_name,
VpcSecurityGroupIds=[security_group_id],
BackupRetentionPeriod=backup_retention_days,
)["DBCluster"]
endpoint = db_result["Endpoint"]
reader_endpoint = db_result["ReaderEndpoint"]
port = db_result["Port"]
arn = db_result["DBClusterArn"]
for i in range(total_instances):
instance_id = f"{cluster_id}-Instance-{i}"
logger.info(
"Creating instances on the cluster",
extra=dict(
cluster_id=cluster_id,
instance_id=instance_id,
node_type=node_type,
),
)
instance_result = self.rds.create_db_instance(
DBInstanceIdentifier=instance_id,
DBClusterIdentifier=cluster_id,
DBInstanceClass=node_type,
Engine=self.db_engine,
PubliclyAccessible=public,
Tags=[{"Key": k, "Value": v} for k, v in (tags or {}).items()],
)
tenant_database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
db = DbInstance(
cluster_id=cluster_id,
arn=arn,
endpoint=endpoint,
reader_endpoint=reader_endpoint,
port=port,
database_name=tenant_database_name,
master_user=master_creds,
)
print(db) # TODO: remove
logger.info("Waiting until DB instances are online. This will take a while (~5 min)")
logger.info("Connecting DB instance using master credentials", extra=dict(endpoint=db.endpoint))
self._check_connection(
connection=db,
credentials=master_creds,
)
tenant_creds = Credentials.new()
logger.info(
"Connection successful. Creating a new database and user",
extra=dict(
endpoint=endpoint,
username=tenant_creds.username,
database_name=tenant_database_name,
),
)
self.create_tenant(
connection=db,
master_user=master_creds,
database_name=tenant_database_name,
tenant_user=tenant_creds,
)
db.user = tenant_creds
return db
@retry(timeout=10 * 60)
def _check_connection(self, connection: DbConnection, credentials: Credentials, database_name: str = "postgres"):
psycopg2.connect(
host=connection.endpoint,
port=connection.port,
dbname=database_name,
user=credentials.username,
password=credentials.password,
connect_timeout=15,
).close()
def _find_instance(self, endpoint: str) -> typing.Optional[DbInstance]:
func = functools.partial(
self.rds.describe_db_clusters, Filters=[{"Name": "engine", "Values": [self.db_engine]}], MaxRecords=100
)
cluster = None
for it in paginate_by_marker(func, "DBClusters"):
if it["Endpoint"] == endpoint:
cluster = it
break
if cluster:
return DbInstance(
endpoint=cluster["Endpoint"],
port=cluster["Port"],
cluster_id=cluster["DBClusterIdentifier"],
arn=cluster["DBClusterArn"],
reader_endpoint=cluster["ReaderEndpoint"],
)
return None
def create_tenant(
self,
connection: DbConnection,
master_user: Credentials,
tenant_user: Credentials,
database_name: typing.Optional[str] = None,
) -> DbInstance:
if not database_name:
database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
try:
con = psycopg2.connect(
host=connection.endpoint,
port=connection.port,
dbname=self.default_master_database_name,
user=master_user.username,
password=master_user.password,
connect_timeout=30,
)
con.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
cur = con.cursor()
cur.execute(
sql.SQL("CREATE DATABASE {}").format(
sql.Identifier(database_name),
),
)
cur.execute(
sql.SQL("CREATE USER {} WITH ENCRYPTED PASSWORD {}").format(
sql.Identifier(tenant_user.username),
sql.Placeholder(),
),
[tenant_user.password],
)
cur.execute(
sql.SQL("GRANT ALL PRIVILEGES ON DATABASE {} TO {}").format(
sql.Identifier(database_name),
sql.Identifier(tenant_user.username),
)
)
# `with` block would create an implicit transaction
con.close()
db = self._find_instance(connection.endpoint)
db.database_name = database_name
db.user = tenant_user
db.master_user = master_user
return db
except psycopg2.Error:
logging.error(
f"Failed to create database and user",
exc_info=True,
extra=dict(endpoint=connection.endpoint, database_name=database_name),
)
raise
class MysqlService(PostgresqlService):
db_type = "MySQL"
db_engine = "aurora-mysql"
db_engine_version = "5.7.12"
default_master_database_name = "mysql"
cluster_id_format = "mysql-{}" # must start with a letter
@retry(timeout=10 * 60)
def _check_connection(self, connection: DbConnection, credentials: Credentials, database_name: str = "postgres"):
pymysql.connect(
host=connection.endpoint,
port=connection.port,
user=credentials.username,
password=credentials.password,
connect_timeout=15,
).close()
def create_tenant(
self,
connection: DbConnection,
master_user: Credentials,
tenant_user: Credentials,
database_name: typing.Optional[str] = None,
) -> DbInstance:
if not database_name:
database_name = self.tenant_database_name_format.format(uuid.uuid4().hex)
try:
with pymysql.connect(
host=connection.endpoint,
port=connection.port,
user=master_user.username,
password=master_user.password,
connect_timeout=15,
) as con:
cur = con.cursor()
cur.execute(f"CREATE DATABASE {database_name}")
cur.execute(f"CREATE USER {tenant_user.username} IDENTIFIED BY '{tenant_user.password}'")
cur.execute(
f"GRANT ALL PRIVILEGES ON {database_name}.* TO {tenant_user.username}@'%' IDENTIFIED BY '{tenant_user.password}'"
)
db = self._find_instance(endpoint=connection.endpoint)
db.database_name = database_name
db.master_user = master_user
db.user = tenant_user
return db
except pymysql.Error:
logger.error(
"Failed to create database and user",
exc_info=True,
extra=dict(endpoint=connection.endpoint, database_name=database_name),
)
raise
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, format=f"%(asctime)s: {logging.BASIC_FORMAT}")
logging.getLogger("botocore").setLevel(logging.INFO)
# m = MysqlService("eu-central-1")
# db = m.create_db(
# total_instances=1,
# db_subnet_group_name="mydbsubnet",
# security_group_id="sg-c11a77ac",
# node_type="db.t3.medium",
# tags={"CostCenter": "hello"},
# public=True,
# )
# print(db)
# p = PostgresqlService("eu-central-1")
# db = p.create_db(
# total_instances=1,
# db_subnet_group_name="mydbsubnet",
# security_group_id="sg-c11a77ac",
# node_type="db.t3.medium",
# tags={"CostCenter": "hello"},
# public=True,
# )
# p.create_tenant(
# connection=DbConnection(
# endpoint="pg0a085614d5eb4f34affd2a24b1fb18c4.cluster-c8v5rp0ouaey.eu-central-1.rds.amazonaws.com",
# port=5432,
# ),
# master_user=Credentials(
# username="u7d0f9621f5a5419288a4ef7ce9fd5fe", password="ca9473879b794dd3a0b0f0d2a8d4aebb"
# ),
# tenant_user=Credentials.new(),
# database_name="mydb",
# )
-107
View File
@@ -1,107 +0,0 @@
import logging
import typing
import boto3
import botocore.client
def paginate_by_marker(
func: callable,
list_field: str,
) -> typing.Iterable:
"""
Simplifies paginating over AWS results with a simpler interface than Paginators.
:param func: A function that accepts a parameter named `Marker`
:param list_field: Result field to iterate on
"""
marker = ""
while True:
result = func(Marker=marker)
items = result.get(list_field, [])
yield from items
if "Marker" not in result:
break
marker = result["Marker"]
logger = logging.getLogger(__name__)
class Rediser:
@classmethod
def new(cls):
return cls(boto3.client("elasticache", region_name='eu-central-1'))
def __init__(self, elasticache: botocore.client.BaseClient):
self.elasticache = elasticache
def _get_redis_acl(
self,
cluster_id: str = None,
status: str = "available",
) -> typing.Optional[dict]:
"""
Finds a Redis replication group by its ID
:param cluster_id: Replication group ID
:param status: Replication status. One of available/starting/modifying/deleting.
:return: RedisInstance or throws LookupError error
"""
try:
redis = self.elasticache.describe_replication_groups(ReplicationGroupId=cluster_id)["ReplicationGroups"][0]
if redis["Status"] != status:
raise Exception(f"Found redis with {cluster_id=} but its status={redis['Status']}")
except self.elasticache.exceptions.ReplicationGroupNotFoundFault:
return None
return dict(
cluster_id=redis["ReplicationGroupId"],
host=redis["NodeGroups"][0]["PrimaryEndpoint"]["Address"],
port=redis["NodeGroups"][0]["PrimaryEndpoint"]["Port"],
arn=redis["ARN"],
user_group_ids=redis["UserGroupIds"],
)
def create_redis_acl(
self,
*,
cluster_id: str,
cluster_description: typing.Optional[str] = None,
node_type: str,
subnet_group_name: str,
security_group_id: str,
):
redis = self._get_redis_acl(cluster_id=cluster_id)
if not redis:
res = self.elasticache.create_replication_group(
ReplicationGroupId=cluster_id,
ReplicationGroupDescription=cluster_description or cluster_id,
ReplicasPerNodeGroup=0,
Engine="redis",
EngineVersion="6.x",
CacheNodeType=node_type,
CacheSubnetGroupName=subnet_group_name,
TransitEncryptionEnabled=True, # must be True to use Redis ACL
UserGroupIds=[
user_group_id,
],
SecurityGroupIds=[
security_group_id,
],
Tags=[{"Key": k, "Value": v} for k, v in tags.items()],
MultiAZEnabled=False,
AutomaticFailoverEnabled=False,
)
def main():
r = Rediser.new()
res = r.create_redis_acl(cluster_id='hey', node_type='cache.t2.micro')
if __name__ == '__main__':
main()
-202
View File
@@ -1,202 +0,0 @@
import concurrent.futures
import functools
import logging
import pprint
import time
import typing
import boto3
import botocore.client
DnsRecord = typing.TypedDict(
"DnsRecord",
{
"purpose": str,
"type": str,
"name": str,
"value": str,
"priority": typing.Optional[str],
},
)
VerificationStatus = typing.Literal["Pending", "Success", "Failed", "TemporaryFailure", "NotStarted"]
logger = logging.getLogger(__name__)
def retry(tries: int = 5, delay: int = 1.5):
def wrapper(fn):
@functools.wraps(fn)
def wrapped(*args, **kwargs):
nonlocal tries
while tries:
try:
return fn(*args, **kwargs)
except Exception:
logger.exception(f'{fn} returned error, retrying {tries} more times')
time.sleep(delay)
tries -= 1
return wrapped
return wrapper
class AwsSesService:
def __init__(self, client: typing.Optional[botocore.client.BaseClient] = None, timeout: int = 5):
self.ses = client or boto3.client("ses")
self.region = self.ses._client_config.region_name
self.timeout = timeout
# @retry()
def configure_domain_identity(
self, domain: str, mail_subdomain: typing.Optional[str] = None
) -> typing.List[DnsRecord]:
"""
Creates up mail identities for sending email through a domain.
DKIM-signed messages help receiving mail servers validate that a message was not forged or altered in transit.
:param domain: domain name to set up email identity
:return: DNS records that need to be set
"""
res_id = self.ses.verify_domain_identity(Domain=domain)
id_token = res_id["VerificationToken"]
res_dkim = self.ses.verify_domain_dkim(Domain=domain)
dns_records = [
{
"purpose": "identity",
"type": "TXT",
"name": domain,
"value": id_token,
},
*[
{
"purpose": "dkim",
"type": "CNAME",
"name": f"{dkim_token}._domainkey.{domain}",
"value": f"{dkim_token}.dkim.amazonses.com",
}
for dkim_token in res_dkim["DkimTokens"]
],
]
if mail_subdomain:
more_dns = self.configure_from_domain(domain, mail_subdomain)
dns_records.extend(more_dns)
return dns_records
@retry()
def check_dkim_verification_status(self, domain: str) -> VerificationStatus:
"""
Returns the DKIM verification status for an email identity.
:param domain: domain used to set up the mail identity
:return: one of {Pending, Success, Failed, TemporaryFailure, NotStarted}
"""
res = self.ses.get_identity_dkim_attributes(Identities=[domain])
return res["DkimAttributes"][domain]["DkimVerificationStatus"]
@retry()
def check_id_verification_status(self, domain: str) -> VerificationStatus:
"""
Returns the id verification status of an email identity.
:param domain: domain used to set up the mail identity
:return: one of {Pending, Success, Failed, TemporaryFailure, NotStarted}
"""
res = self.ses.get_identity_verification_attributes(Identities=[domain])
return res["VerificationAttributes"][domain]["VerificationStatus"]
def configure_from_domain(self, domain: str, mail_subdomain: str) -> typing.List[DnsRecord]:
"""
Allows using `mail_subdomain` as FROM address when sending emails.
Messages sent through Amazon SES will be marked as originating from your domain instead of a subdomain of amazon.com.
:param domain: mail identity
:param mail_subdomain: a subdomain
:return: DNS records that need to be set
"""
_ = self.ses.set_identity_mail_from_domain(
Identity=domain,
MailFromDomain=mail_subdomain,
BehaviorOnMXFailure="UseDefaultValue", # uses $region.amazonses.com if MX records not present on the mail subdomain
)
return [
{
"purpose": "mail_from_domain",
"type": "MX",
"name": mail_subdomain,
"value": f"feedback-smtp.{self.region}.amazonses.com",
"priority": "10",
},
{
"purpose": "mail_from_domain",
"type": "TXT",
"name": mail_subdomain,
"value": f'"v=spf1 include:amazonses.com ~all"',
},
]
@retry()
def check_custom_from_domain_status(self, domain: str) -> VerificationStatus:
"""
Returns the status of custom email FROM request verification.
:param domain: domain used to set up email identity
:return: one of {Pending, Success, Failed, TemporaryFailure}
"""
res = self.ses.get_identity_mail_from_domain_attributes(Identities=[domain])
return res["MailFromDomainAttributes"][domain].get("MailFromDomainStatus", 'NotStarted')
def bulk_check_domain_verification(self, identity_domain: str, timeout: int = 5) -> dict:
tasks = {
"status_id": functools.partial(self.check_id_verification_status, identity_domain),
"status_dkim": functools.partial(self.check_dkim_verification_status, identity_domain),
"status_custom": functools.partial(self.check_custom_from_domain_status, identity_domain),
}
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as pool:
futures = [pool.submit(t) for _, t in tasks.items()]
results = dict(zip(tasks.keys(), [f.result(timeout=timeout) for f in futures]))
return {
**results,
"identity_domain": identity_domain,
}
@retry()
def bulk_configure_email_domains(
self,
identity_domain: str,
mail_subdomain: typing.Optional[str] = None,
timeout: int = 10,
) -> dict:
tasks = {
"dns": functools.partial(self.configure_domain_identity, identity_domain, mail_subdomain),
}
# send all requests in parallel
with concurrent.futures.ThreadPoolExecutor(max_workers=10) as pool:
futures = [pool.submit(t) for _, t in tasks.items()]
results = dict(zip(tasks.keys(), [f.result(timeout=timeout) for f in futures]))
return {
**results,
"identity_domain": identity_domain,
"mail_subdomain": mail_subdomain,
}
if __name__ == "__main__":
ses = AwsSesService()
logging.basicConfig(level=logging.INFO)
# print(ses.bulk_configure_email_domains('abdus100.dev', 'mail.abdus100.dev'))
# print(ses.check_custom_from_domain_status('abdus101.dev'))
tasks = [
functools.partial(ses.bulk_check_domain_verification, d) for d in ['abdus.dev', 'abdus1.dev', 'abdus2.dev']
] * 3
with concurrent.futures.ThreadPoolExecutor(max_workers=20) as pool:
futures = [pool.submit(t) for t in tasks]
results = [f.result(timeout=5) for f in futures]
pprint.pp(results)
-133
View File
@@ -1,133 +0,0 @@
from unittest import mock
import pytest
from aws_ses import AwsSesService
# fmt: off
_response_verify_domain_dkim = {'DkimTokens': ['token1', 'token2', 'token3'], 'ResponseMetadata': {'RequestId': '6e4da72e-3063-41da-9434-dae7ef19c78e', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:27:40 GMT', 'content-type': 'text/xml', 'content-length': '469', 'connection': 'keep-alive', 'x-amzn-requestid': '6e4da72e-3063-41da-9434-dae7ef19c78e'}, 'RetryAttempts': 0}}
_response_verify_domain_identity = {'VerificationToken': 'token', 'ResponseMetadata': {'RequestId': 'b733069c-115f-4457-b6cb-14e44f29d17f', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:27:12 GMT', 'content-type': 'text/xml', 'content-length': '370', 'connection': 'keep-alive', 'x-amzn-requestid': 'b733069c-115f-4457-b6cb-14e44f29d17f'}, 'RetryAttempts': 0}}
_response_get_identity_dkim_attributes = {'DkimAttributes': {'example.com': {'DkimEnabled': True, 'DkimVerificationStatus': 'Success', 'DkimTokens': ['bgobmdqs2vbxr6ygwbmg753knfu6tos7', 'm6o6w6twpmwzwe3igpxxo7qsvwp3hlgj', 'zqfz4ekdelp7m2u2phsyad3h5dsfle4v']}}, 'ResponseMetadata': {'RequestId': 'ba466be1-79bc-499b-9766-ca1acd7d1c08', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 16:49:53 GMT', 'content-type': 'text/xml', 'content-length': '833', 'connection': 'keep-alive', 'x-amzn-requestid': 'ba466be1-79bc-499b-9766-ca1acd7d1c08'}, 'RetryAttempts': 0}}
_response_get_identity_dkim_attributes_pending = {'DkimAttributes': {'example.com': {'DkimEnabled': True, 'DkimVerificationStatus': 'Pending', 'DkimTokens': ['bgobmdqs2vbxr6ygwbmg753knfu6tos7', 'm6o6w6twpmwzwe3igpxxo7qsvwp3hlgj', 'zqfz4ekdelp7m2u2phsyad3h5dsfle4v']}}, 'ResponseMetadata': {'RequestId': 'b24732d9-01e1-4266-8b34-550a9ecf634d', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:30:03 GMT', 'content-type': 'text/xml', 'content-length': '833', 'connection': 'keep-alive', 'x-amzn-requestid': 'b24732d9-01e1-4266-8b34-550a9ecf634d'}, 'RetryAttempts': 0}}
_response_get_identity_verification_attributes = {'VerificationAttributes': {'example.com': {'VerificationStatus': 'Success', 'VerificationToken': 'XzMvLQqOPdC0iJ02YTdjqWBJr7FTqhGtAsEcfpUwfoA='}}, 'ResponseMetadata': {'RequestId': '022c3e6e-a6e8-426f-bf4d-149f71ca7913', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 16:56:19 GMT', 'content-type': 'text/xml', 'content-length': '637', 'connection': 'keep-alive', 'x-amzn-requestid': '022c3e6e-a6e8-426f-bf4d-149f71ca7913'}, 'RetryAttempts': 0}}
_response_get_identity_verification_attributes_pending = {'VerificationAttributes': {'example.com': {'VerificationStatus': 'Pending', 'VerificationToken': 'XzMvLQqOPdC0iJ02YTdjqWBJr7FTqhGtAsEcfpUwfoA='}}, 'ResponseMetadata': {'RequestId': '022c3e6e-a6e8-426f-bf4d-149f71ca7913', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 16:56:19 GMT', 'content-type': 'text/xml', 'content-length': '637', 'connection': 'keep-alive', 'x-amzn-requestid': '022c3e6e-a6e8-426f-bf4d-149f71ca7913'}, 'RetryAttempts': 0}}
_response_set_identity_mail_from_domain = {'ResponseMetadata': {'RequestId': 'c04acd33-40be-46cd-9960-c9e51397dc48', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:03:50 GMT', 'content-type': 'text/xml', 'content-length': '266', 'connection': 'keep-alive', 'x-amzn-requestid': 'c04acd33-40be-46cd-9960-c9e51397dc48'}, 'RetryAttempts': 0}}
_response_get_identity_mail_from_domain_attributes = {'MailFromDomainAttributes': {'example.com': {'MailFromDomain': 'email.abdus.dev', 'MailFromDomainStatus': 'Success', 'BehaviorOnMXFailure': 'UseDefaultValue'}}, 'ResponseMetadata': {'RequestId': '63049966-c062-47f4-8238-145bd533d880', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:37:06 GMT', 'content-type': 'text/xml', 'content-length': '687', 'connection': 'keep-alive', 'x-amzn-requestid': '63049966-c062-47f4-8238-145bd533d880'}, 'RetryAttempts': 0}}
_response_get_identity_mail_from_domain_attributes_pending = {'MailFromDomainAttributes': {'example.com': {'MailFromDomain': 'email.abdus.dev', 'MailFromDomainStatus': 'Pending', 'BehaviorOnMXFailure': 'UseDefaultValue'}}, 'ResponseMetadata': {'RequestId': '63049966-c062-47f4-8238-145bd533d880', 'HTTPStatusCode': 200, 'HTTPHeaders': {'date': 'Mon, 03 Jan 2022 17:37:06 GMT', 'content-type': 'text/xml', 'content-length': '687', 'connection': 'keep-alive', 'x-amzn-requestid': '63049966-c062-47f4-8238-145bd533d880'}, 'RetryAttempts': 0}}
# fmt: on
@pytest.fixture()
def mock_ses() -> mock.MagicMock:
with mock.patch("boto3.client") as m:
m_ses = m.return_value
yield m_ses
def test_init_domain_verification(mock_ses):
m_verify_id = mock_ses.verify_domain_identity
m_verify_dkim = mock_ses.verify_domain_dkim
m_verify_id.return_value = _response_verify_domain_identity
m_verify_dkim.return_value = _response_verify_domain_dkim
domain = "example.com"
dns = AwsSesService(mock_ses).configure_domain_identity(domain)
m_verify_id.assert_called_once_with(Domain=domain)
m_verify_dkim.assert_called_once_with(Domain=domain)
assert dns == [
{
"type": "TXT",
"name": domain,
"value": "token",
},
{
"type": "CNAME",
"name": f"token1._domainkey.{domain}",
"value": f"token1.dkim.amazonses.com",
},
{
"type": "CNAME",
"name": f"token2._domainkey.{domain}",
"value": f"token2.dkim.amazonses.com",
},
{
"type": "CNAME",
"name": f"token3._domainkey.{domain}",
"value": f"token3.dkim.amazonses.com",
},
]
def test_get_dkim_status(mock_ses):
m = mock_ses.get_identity_dkim_attributes
m.return_value = _response_get_identity_dkim_attributes_pending
domain = "example.com"
status = AwsSesService(mock_ses).check_dkim_verification_status(domain)
assert status == "Pending"
m.assert_called_once_with(Identities=[domain])
# ===
m.reset_mock()
m.return_value = _response_get_identity_dkim_attributes
status = AwsSesService(mock_ses).check_dkim_verification_status(domain)
assert status == "Success"
def test_get_id_status(mock_ses):
m = mock_ses.get_identity_verification_attributes
m.return_value = _response_get_identity_verification_attributes
domain = "example.com"
status = AwsSesService(mock_ses).check_id_verification_status(domain)
assert status == "Success"
m.assert_called_once_with(Identities=[domain])
def test_set_custom_from_domain(mock_ses):
mock_ses._client_config.region_name = "fake-region"
m = mock_ses.set_identity_mail_from_domain
m.return_value = _response_set_identity_mail_from_domain
dns = AwsSesService(mock_ses).configure_from_domain("example.com", "mail.example.com")
assert dns == [
{
"type": "MX",
"name": "mail.example.com",
"value": f"feedback-smtp.fake-region.amazonses.com",
"priority": "10",
},
{"type": "TXT", "name": "mail.example.com", "value": f'"v=spf1 include:amazonses.com ~all"'},
]
m.assert_called_once_with(
Identity="example.com",
MailFromDomain="mail.example.com",
BehaviorOnMXFailure="UseDefaultValue",
)
def test_get_mail_from_status(mock_ses):
m = mock_ses.get_identity_mail_from_domain_attributes
m.return_value = _response_get_identity_mail_from_domain_attributes_pending
domain = "example.com"
status = AwsSesService(mock_ses).check_custom_from_domain_status(domain)
assert status == "Pending"
m.assert_called_once_with(Identities=[domain])
# ===
m.reset_mock()
m.return_value = _response_get_identity_mail_from_domain_attributes
status = AwsSesService(mock_ses).check_custom_from_domain_status(domain)
assert status == "Success"
m.assert_called_once_with(Identities=[domain])
-237
View File
@@ -1,237 +0,0 @@
#!/usr/bin/env python3.9
import base64
import json
import logging
import pprint
import subprocess
import time
import webbrowser
from pathlib import Path
import httpx
import typer
hetzner_token = 'zdlkmAF9EswvOIEGvWgutULYTmUpa3c44YcKZOB0jymQfgvra7kRVJxuLiFTouM9'
client = httpx.Client(base_url='https://api.hetzner.cloud/v1', headers={'Authorization': f'Bearer {hetzner_token}'})
datacenters = [
{'id': 2, 'name': 'nbg1-dc3', 'description': 'Nuremberg 1 DC 3'},
{'id': 3, 'name': 'hel1-dc2', 'description': 'Helsinki 1 DC 2'},
{'id': 4, 'name': 'fsn1-dc14', 'description': 'Falkenstein 1 DC14'},
{'id': 5, 'name': 'ash-dc1', 'description': 'Ashburn DC1'},
]
server_types = {
'cx11': {'cores': 1, 'memory': 2.0, 'disk': 20},
'cx21': {'cores': 2, 'memory': 4.0, 'disk': 40},
'cx31': {'cores': 2, 'memory': 8.0, 'disk': 80},
'cx41': {'cores': 4, 'memory': 16.0, 'disk': 160},
'cx51': {'cores': 8, 'memory': 32.0, 'disk': 240},
'ccx11': {'cores': 2, 'memory': 8.0, 'disk': 80},
'ccx21': {'cores': 4, 'memory': 16.0, 'disk': 160},
'ccx31': {'cores': 8, 'memory': 32.0, 'disk': 240},
'ccx41': {'cores': 16, 'memory': 64.0, 'disk': 360},
'ccx51': {'cores': 32, 'memory': 128.0, 'disk': 600},
'cpx11': {'cores': 2, 'memory': 2.0, 'disk': 40},
'cpx21': {'cores': 3, 'memory': 4.0, 'disk': 80},
'cpx31': {'cores': 4, 'memory': 8.0, 'disk': 160},
'cpx41': {'cores': 8, 'memory': 16.0, 'disk': 240},
'cpx51': {'cores': 16, 'memory': 32.0, 'disk': 360},
'ccx12': {'cores': 2, 'memory': 8.0, 'disk': 80},
'ccx22': {'cores': 4, 'memory': 16.0, 'disk': 160},
'ccx32': {'cores': 8, 'memory': 32.0, 'disk': 240},
'ccx42': {'cores': 16, 'memory': 64.0, 'disk': 360},
'ccx52': {'cores': 32, 'memory': 128.0, 'disk': 600},
'ccx62': {'cores': 48, 'memory': 192.0, 'disk': 960},
}
def create_server(name: str, datacenter: str):
payload = {
"datacenter": datacenter,
"image": "ubuntu-20.04",
"name": name,
"server_type": "cpx21",
"ssh_keys": ["entropy"],
}
res = client.post('/servers', json=payload)
res.raise_for_status()
data = res.json()
server = data['server']
pprint.pprint(server)
return {
'id': server['id'],
'name': server['name'],
'ip': server['public_net']['ipv4']['ip'],
}
def find_server(name: str) -> dict:
res = client.get('/servers')
res.raise_for_status()
data = res.json()
server = [it for it in data['servers'] if it['name'] == name][0]
return {
'id': server['id'],
'name': server['name'],
'ip': server['public_net']['ipv4']['ip'],
}
def delete_server(name: str) -> None:
server = find_server(name)
logging.info(f'deleting server={name} with id={server["id"]}')
res = client.delete(f'/servers/{server["id"]}')
res.raise_for_status()
logging.info(f'deleted server={name}')
def copy_to_clipboard(text: str) -> None:
subprocess.Popen(['pbcopy'], stdin=subprocess.PIPE).communicate(input=text.encode())
class Server:
ssh_options = ['-o', 'StrictHostKeyChecking=no']
def __init__(self, ip_address: str, name: str = None):
self.ip = ip_address
self.ssh_host = f'root@{self.ip}'
self.name = name
@classmethod
def from_hetzner(cls, server: dict):
return cls(server['ip'], name=server['name'])
def run_command(self, cmd: str, return_output: bool = False) -> str:
logging.debug(f'running command={cmd} on {self.ssh_host}')
extra_args = {}
if return_output:
extra_args['stdout'] = subprocess.PIPE
p = subprocess.run(['ssh', *self.ssh_options, self.ssh_host, '--', cmd], text=True, **extra_args)
p.check_returncode()
return p.stdout
def install_package(self, packages: str) -> str:
return self.run_command(f'apt-get install -y {packages}')
def install_shell(self):
self.install_package('fish')
self.run_command('chsh -s /usr/bin/fish')
def install_microk8s(self):
self.install_package('snapd')
self.run_command('snap install microk8s --classic')
self.run_command('ufw allow in on cni0 && ufw allow out on cni0 && ufw default allow routed')
self.run_command('microk8s enable dashboard dns rbac ingress storage prometheus registry')
def get_microk8s_token(self):
raw_json = self.run_command('microk8s kubectl -n kube-system get secret -o json', return_output=True)
secrets = json.loads(raw_json)['items']
default_token_secret = [it for it in secrets if it['metadata']['name'].startswith('default-token')][0]
token_base64 = default_token_secret['data']['token']
token = base64.decodebytes(token_base64.encode()).decode()
return token
def show_microk8s_access_info(self):
logging.info('getting default service account token')
token = self.get_microk8s_token()
print(token)
logging.info('copying token to clipboard')
copy_to_clipboard(token)
logging.info('getting dashboard ip')
proxy = self.get_microk8s_dashboard_proxy()
print(proxy['ssh'])
copy_to_clipboard(proxy['ssh'])
logging.info('opening url in the browser')
print(proxy['url'])
webbrowser.open_new_tab(proxy['url'])
def get_microk8s_dashboard_proxy(self):
cluster_ip = self.run_command(
'microk8s kubectl get -n kube-system -o json service/kubernetes-dashboard -o jsonpath="{.spec.clusterIP}"',
return_output=True,
).strip()
port, local_port = 443, 18433
return {
'ssh': f'ssh -vNL {local_port}:{cluster_ip}:{port} {self.ssh_host}',
'url': f'https://localhost:{local_port}',
}
def write_kubeconfig(self, filename: str):
logging.info(f'writing kubeconfig to {filename}')
kubeconfig = self.run_command('microk8s config', return_output=True)
file_path = Path(filename).expanduser()
file_path.parent.mkdir(parents=True, exist_ok=True)
file_path.write_text(kubeconfig)
def install_micro_editor(self):
logging.info('installing micro editor')
self.run_command('curl https://getmic.ro | bash && mv ./micro /usr/bin/')
def setup(self):
self.run_command('apt-get update')
self.install_shell()
self.install_micro_editor()
self.install_microk8s()
logging.info('waiting until pods are online')
time.sleep(10)
self.show_microk8s_access_info()
self.write_kubeconfig(f'~/Desktop/kubeconfig-{self.name}.yml')
cli = typer.Typer(invoke_without_command=False)
@cli.command()
def create(
name: str,
datacenter: str = typer.Option('nbg1-dc3'),
):
create_server(name=name, datacenter=datacenter)
@cli.command()
def destroy(
name: str,
):
delete_server(name=name)
@cli.command()
def setup(
name: str,
):
server = find_server(name)
server = Server.from_hetzner(server)
server.setup()
@cli.command()
def show_k8s(
name: str,
):
server = Server.from_hetzner(find_server(name))
server.show_microk8s_access_info()
server.write_kubeconfig(f'~/Desktop/kubeconfig-{server.name}.yml')
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO)
cli()
-72
View File
@@ -1,72 +0,0 @@
import contextlib
import dataclasses
import plistlib
import sqlite3
import subprocess
from pathlib import Path
from typing import Optional
DB_PATH = Path("~/Library/Application Support/Dock/desktoppicture.db").expanduser()
@dataclasses.dataclass
class SpaceInfo:
uuid: str
image_path: Optional[Path]
data_id: Optional[Path]
@contextlib.contextmanager
def get_connection() -> sqlite3.Connection:
con = sqlite3.connect(str(DB_PATH))
con.row_factory = sqlite3.Row
yield con
con.close()
def get_current_wallpapers() -> list[SpaceInfo]:
sql = """
select distinct s.space_uuid, pr.key, d.rowid, d.value from preferences pr
left join payload d on d.rowid = pr.data_id
left join pictures p on p.rowid = pr.picture_id
inner join spaces s on s.rowid = p.space_id
where key = 1
order by space_uuid, key
"""
spaces = list_spaces()
with get_connection() as con:
cur = con.execute(sql)
space_props = {}
for i, row in enumerate(cur):
space_uuid = row["space_uuid"]
data_id = row["rowid"]
image_path = row["value"]
if image_path:
image_path = Path(image_path).expanduser()
space_props[space_uuid] = SpaceInfo(uuid=space_uuid, image_path=image_path, data_id=data_id)
return [space_props.get(id, SpaceInfo(uuid="", image_path=None, data_id=None)) for id in spaces]
def read_plist(plist_path):
with plist_path.open("rb") as f:
return plistlib.load(f)
def list_spaces():
spaces_plist_path = Path("~/Library/Preferences/com.apple.spaces.plist").expanduser()
data = read_plist(spaces_plist_path)
spaces = data["SpacesDisplayConfiguration"]["Management Data"]["Monitors"][0]["Spaces"]
return [s["uuid"] for s in spaces]
def set_wallpaper(image_path: Path):
script = f'tell application "Finder" to set desktop picture to "{image_path.expanduser().resolve()}" as POSIX file'
subprocess.run(["osascript", "-e", script], check=True)
if __name__ == "__main__":
pass
-523
View File
@@ -1,523 +0,0 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Prompter</title>
<!-- Persist plugin must load before Alpine core -->
<script defer src="https://cdn.jsdelivr.net/npm/@alpinejs/persist@3.x.x/dist/cdn.min.js"></script>
<script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
<style>
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
body {
display: flex;
height: 100vh;
overflow: hidden;
font-family: system-ui, -apple-system, sans-serif;
background: #0f0f0f;
color: #e0e0e0;
}
/* ── Image panel (left 50%) ──────────────────────────────────────────── */
.image-panel {
width: 50%;
height: 100vh;
display: flex;
flex-direction: column;
align-items: center;
justify-content: center;
background: #111;
padding: 20px;
gap: 10px;
border-right: 1px solid #1e1e1e;
}
.image-panel img {
max-width: 100%;
max-height: calc(100vh - 72px);
object-fit: contain;
}
.image-meta {
display: flex;
flex-direction: column;
align-items: center;
gap: 4px;
}
.image-filename {
font-size: 12px;
color: #666;
word-break: break-all;
text-align: center;
max-width: 100%;
}
.image-counter {
display: flex;
align-items: center;
gap: 4px;
font-size: 12px;
color: #555;
}
.image-counter input[type="number"] {
background: transparent;
border: none;
border-bottom: 1px solid #444;
color: #999;
font-family: inherit;
font-size: 12px;
outline: none;
padding: 0 2px;
text-align: right;
width: 3.5ch;
-moz-appearance: textfield;
}
.image-counter input[type="number"]:focus {
border-bottom-color: #4a6cf7;
color: #e0e0e0;
}
.image-counter input::-webkit-outer-spin-button,
.image-counter input::-webkit-inner-spin-button { -webkit-appearance: none; }
/* ── Options panel (right 50%) ───────────────────────────────────────── */
.options-panel {
width: 50%;
height: 100vh;
display: flex;
flex-direction: column;
overflow: hidden;
}
.panel-top,
.panel-bottom {
flex: 1 1 0;
min-height: 0;
overflow-y: auto;
padding: 24px;
display: flex;
flex-direction: column;
gap: 8px;
}
.panel-bottom {
border-top: 1px solid #1e1e1e;
}
.section-title {
font-size: 11px;
color: #555;
text-transform: uppercase;
letter-spacing: 0.08em;
margin-bottom: 4px;
}
.option {
border: 1px solid #2a2a2a;
border-radius: 8px;
padding: 12px 14px;
cursor: pointer;
transition: border-color 0.1s;
}
.option:hover { border-color: #3a3a3a; }
.option.active { border-color: #4a6cf7; background: rgba(74, 108, 247, 0.06); }
.option-meta {
display: flex;
align-items: center;
gap: 8px;
margin-bottom: 10px;
font-size: 12px;
color: #666;
}
.saved-badge { color: #4caf50; margin-left: auto; font-size: 11px; }
.kbd {
display: inline-block;
background: #1e1e1e;
border: 1px solid #333;
border-radius: 3px;
padding: 1px 6px;
font-size: 11px;
font-family: monospace;
color: #888;
white-space: nowrap;
}
textarea {
width: 100%;
background: #1a1a1a;
border: 1px solid #333;
border-radius: 6px;
color: #e0e0e0;
font-size: 14px;
line-height: 1.5;
padding: 10px;
resize: vertical;
min-height: 90px;
font-family: inherit;
outline: none;
transition: border-color 0.1s;
}
textarea:focus { border-color: #4a6cf7; }
.btn-save {
margin-top: 10px;
padding: 8px 18px;
background: #4a6cf7;
color: #fff;
border: none;
border-radius: 6px;
font-size: 14px;
cursor: pointer;
transition: background 0.1s;
}
.btn-save:hover { background: #3b5de8; }
.history-section { margin-top: 8px; }
.history-item { margin-top: 8px; }
.history-text {
font-size: 13px;
color: #bbb;
white-space: pre-wrap;
word-break: break-word;
line-height: 1.4;
}
.blocks-grid {
display: grid;
grid-template-columns: 1fr 1fr;
gap: 6px;
margin-top: 4px;
}
.block-pill {
background: #1e1e1e;
border: 1px solid #333;
border-radius: 8px;
color: #ccc;
cursor: pointer;
font-family: inherit;
font-size: 12px;
line-height: 1.4;
padding: 6px 10px;
text-align: left;
transition: border-color 0.1s, background 0.1s;
white-space: pre-wrap;
word-break: break-word;
}
.block-pill:hover {
background: #2a2a2a;
border-color: #4a6cf7;
color: #e0e0e0;
}
</style>
</head>
<body x-data="prompter" @keydown.window="handleKey($event)">
<!-- ── Left: image ───────────────────────────────────────────────────── -->
<div class="image-panel">
<img :src="`/image/${currentIndex}`" :alt="currentImage" />
<div class="image-meta">
<div class="image-filename" x-text="currentImage.split('/').pop()"></div>
<div class="image-counter">
<input
type="number"
min="1"
:max="images.length"
:value="currentIndex + 1"
@change="jumpTo($event.target.valueAsNumber - 1)"
@keydown.stop
@click.stop
/>
<span>/ <span x-text="images.length"></span></span>
</div>
</div>
</div>
<!-- ── Right: options ────────────────────────────────────────────────── -->
<div class="options-panel">
<!-- ── Top half: prompt + history ──────────────────────────────────── -->
<div class="panel-top">
<div class="section-title">Prompt</div>
<!-- Custom prompt -->
<div class="option" :class="{ active: selectedOption === 'custom' }" @click="selectCustom()">
<div class="option-meta">
<span class="kbd">alt+1</span>
<span>Custom prompt</span>
<span class="saved-badge" x-show="selections[currentImage]">&#10003;&nbsp;saved</span>
</div>
<textarea
x-ref="textarea"
x-model="customPrompt"
@focus="selectedOption = 'custom'"
placeholder="Enter a prompt&#8230;"
rows="5"
></textarea>
<button class="btn-save" @click.stop="saveAndNext()">Save &amp; Next</button>
</div>
<!-- History -->
<div class="history-section" x-show="promptHistory.length > 0">
<div class="section-title">Previously used</div>
<template x-for="(prompt, i) in promptHistory" :key="prompt">
<div
class="option history-item"
:class="{ active: selectedOption === i }"
@click="selectHistory(i)"
>
<div class="option-meta">
<span class="kbd" x-text="`alt+${i + 2}`"></span>
</div>
<div class="history-text" x-text="prompt"></div>
</div>
</template>
</div>
</div>
<!-- ── Bottom half: building blocks ────────────────────────────────── -->
<div class="panel-bottom" x-show="buildingBlocks.length > 0">
<div class="section-title">Building blocks</div>
<div class="blocks-grid">
<template x-for="block in buildingBlocks" :key="block">
<button class="block-pill" @click="appendBlock(block)" x-text="block"></button>
</template>
</div>
</div>
</div>
<script>
/**
* @typedef {{ images: string[], canned_prompts: string[], building_blocks: string[], existing_prompts: Record<string, string> }} ApiImages
*/
document.addEventListener('alpine:init', () => {
Alpine.data('prompter', () => ({
/** @type {string[]} Ordered list of absolute image paths from the server */
images: [],
/** @type {number} Zero-based index of the currently displayed image */
currentIndex: 0,
/** @type {Record<string, string>} Absolute image path → saved prompt text */
selections: {},
/** @type {string} Current value of the custom-prompt textarea */
customPrompt: '',
/** @type {'custom' | number | null} Which option row is highlighted */
selectedOption: null,
/** @type {string[]} Reusable text fragments loaded from the server */
buildingBlocks: [],
/**
* Prompt history persisted to localStorage via Alpine Persist.
* Ordered most-recent-first; no duplicates.
* @type {string[]}
*/
promptHistory: Alpine.$persist([]).as('prompter_history'),
/**
* Absolute path of the image currently on screen.
* @returns {string}
*/
get currentImage() {
return this.images[this.currentIndex] ?? '';
},
/**
* Fetch the image list from the server, merge canned prompts into
* history, and restore any previously saved prompt for the first image.
*/
async init() {
/** @type {ApiImages} */
const data = await fetch('/api/images').then(r => r.json());
this.images = data.images;
this.selections = { ...data.existing_prompts };
this.buildingBlocks = data.building_blocks ?? [];
// Append canned prompts not already present in the persisted history.
const inHistory = new Set(this.promptHistory);
for (const p of data.canned_prompts) {
if (!inHistory.has(p)) {
this.promptHistory.push(p);
}
}
// Restore a previously saved prompt for the first image (if any).
const saved = this.selections[this.currentImage];
if (saved) {
this.customPrompt = saved;
this.selectedOption = 'custom';
}
},
/**
* Append a building block to the textarea, separated by a blank line.
* If the textarea is empty the block becomes the entire content.
* @param {string} block
*/
appendBlock(block) {
const cur = this.customPrompt.trimEnd();
this.customPrompt = cur ? cur + '\n' + block : block;
this.selectedOption = 'custom';
this.$nextTick(() => this.$refs.textarea.focus());
},
/** Highlight the custom-prompt option and focus the textarea. */
selectCustom() {
this.selectedOption = 'custom';
this.$nextTick(() => this.$refs.textarea.focus());
},
/**
* Copy a history entry into the textarea and focus it.
* @param {number} i - Index into promptHistory
*/
selectHistory(i) {
this.customPrompt = this.promptHistory[i];
this.selectedOption = 'custom';
this.$nextTick(() => this.$refs.textarea.focus());
},
/**
* POST the current textarea content to the server and update local
* state. Does nothing if the textarea is blank.
*/
async saveCurrentIfNeeded() {
const prompt = this.customPrompt.trim();
if (!prompt) return;
const filePath = this.currentImage;
this.selections[filePath] = prompt;
this._addToHistory(prompt);
await fetch('/api/save', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ file_path: filePath, prompt }),
});
},
/**
* Prepend a prompt to history, removing any earlier duplicate entry.
* @param {string} prompt
*/
_addToHistory(prompt) {
this.promptHistory = [
prompt,
...this.promptHistory.filter(p => p !== prompt),
];
},
/**
* Populate the textarea with the saved prompt for the current image,
* or clear it if no prompt has been saved yet.
*/
_loadCurrentSelection() {
const saved = this.selections[this.currentImage];
this.customPrompt = saved ?? '';
this.selectedOption = saved ? 'custom' : null;
},
/** Save the current prompt (if any) then advance to the next image. */
async navigateNext() {
if (this.currentIndex >= this.images.length - 1) return;
await this.saveCurrentIfNeeded();
this.currentIndex++;
this._loadCurrentSelection();
},
/** Save the current prompt (if any) then go back to the previous image. */
async navigatePrev() {
if (this.currentIndex <= 0) return;
await this.saveCurrentIfNeeded();
this.currentIndex--;
this._loadCurrentSelection();
},
/** Save the current prompt (if any) and advance to the next image. */
async saveAndNext() {
await this.saveCurrentIfNeeded();
if (this.currentIndex < this.images.length - 1) {
this.currentIndex++;
this._loadCurrentSelection();
}
},
/**
* Save the current prompt (if any) then jump to a specific index.
* Out-of-range values are clamped silently.
* @param {number} idx - Zero-based target index
*/
async jumpTo(idx) {
const target = Math.max(0, Math.min(idx, this.images.length - 1));
if (target === this.currentIndex) return;
await this.saveCurrentIfNeeded();
this.currentIndex = target;
this._loadCurrentSelection();
},
/**
* Global keyboard handler (attached to window via @keydown.window).
*
* Shortcuts:
* Alt+1 → focus the custom-prompt textarea
* Alt+N (N≥2) → load history[N-2] into the textarea and focus it
* → / k → next image (k ignored while textarea is focused)
* ← / j → prev image (j ignored while textarea is focused)
*
* Arrow keys navigate from anywhere, including inside the textarea,
* and auto-save the current prompt if non-empty.
*
* @param {KeyboardEvent} e
*/
handleKey(e) {
const inTextarea =
this.$refs.textarea &&
document.activeElement === this.$refs.textarea;
if (inTextarea) return;
if (e.altKey) {
const m = e.code.match(/^Digit(\d)$/);
if (m) {
const n = parseInt(m[1], 10);
e.preventDefault();
if (n === 1) {
this.selectCustom();
} else if (n >= 2 && n - 2 < this.promptHistory.length) {
this.selectHistory(n - 2);
}
return;
}
}
if (e.key === 'ArrowRight' || e.key === 'k') {
e.preventDefault();
this.navigateNext();
} else if (e.key === 'ArrowLeft' || e.key === 'j') {
e.preventDefault();
this.navigatePrev();
}
},
}));
});
</script>
</body>
</html>
-205
View File
@@ -1,205 +0,0 @@
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
# "bottle",
# ]
# ///
"""Browse images one-by-one and assign text prompts via a local web UI.
Saves results to a JSONL file: {"filename": "...", "prompt": "..."}
Images navigated past without a prompt produce no output line.
"""
import argparse
import json
import socket
import sys
import threading
import time
import webbrowser
from pathlib import Path
import bottle
# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
CANNED_PROMPTS: list[str] = [
"photorealistic. skin must be flawless",
]
"""Pre-populated prompt history shown in the UI on first launch."""
BUILDING_BLOCKS: list[str] = [
"make her prettier",
"age her by 5 years and make her look like a 25 year old",
"keep the dark skin",
"wet skin",
"skin covered with baby oil, shiny skin",
"soaking wet skin and hair, wet clothes clinging to her body",
"long, flowy hair",
"japanese",
"greek",
"no extra or missing fingers. each hand must have 5 fingers with the same hand pose as the original image",
# expression
"direct her gaze at the camera",
"give her a serious and seductive look",
"give her a playful expression",
"slightly parted lips",
"french kiss, eyes closed, tongue out",
"closed eyes",
"lower her eyelids and slightly part her lips in a seductive fashion",
"exaggerated expression",
"soft dramatic lighting",
# outfit
"replace the outfit with glossy latex",
"remove the tan lines, she doesn't wear white bra",
]
"""Palette of text snippets the user can append to any prompt."""
VALID_EXTENSIONS: frozenset[str] = frozenset({".webp", ".jpg", ".jpeg", ".png"})
HTML: str = (Path(__file__).parent / "prompter.html").read_text(encoding="utf-8")
# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------
def find_free_port() -> int:
"""Bind to port 0 and return the OS-assigned port number."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return s.getsockname()[1]
class PrompterApp:
def __init__(self, images: list[str], output_path: Path) -> None:
self.images = images
self.output_path = output_path
self.app = bottle.Bottle()
self.app.route("/")(self.index)
self.app.route("/api/images")(self.api_images)
self.app.route("/image/<idx:int>")(self.serve_image)
self.app.route("/api/save", method="POST")(self.api_save)
def index(self) -> str:
return HTML
def api_images(self) -> str:
"""Return image list, canned prompts, building blocks, and existing prompts."""
stem_to_path: dict[str, str] = {Path(p).stem: p for p in self.images}
existing_by_stem = self._load_existing_prompts()
existing_prompts = {
stem_to_path[stem]: prompt
for stem, prompt in existing_by_stem.items()
if stem in stem_to_path
}
bottle.response.content_type = "application/json"
return json.dumps(
{
"images": self.images,
"canned_prompts": CANNED_PROMPTS,
"building_blocks": BUILDING_BLOCKS,
"existing_prompts": existing_prompts,
}
)
def serve_image(self, idx: int) -> bottle.HTTPResponse:
if idx < 0 or idx >= len(self.images):
bottle.abort(404, "Image not found")
path = Path(self.images[idx])
return bottle.static_file(path.name, root=str(path.parent))
def api_save(self) -> str:
data: dict = bottle.request.json or {}
file_path: str = data.get("file_path", "")
prompt: str = data.get("prompt", "")
if file_path and prompt:
self._save_entry(file_path, prompt)
bottle.response.content_type = "application/json"
return json.dumps({"ok": True})
def _load_existing_prompts(self) -> dict[str, str]:
"""Read the output JSONL and return a mapping of stem → prompt.
When a filename appears multiple times the last entry wins, so re-running
the tool and overwriting a previous prompt works naturally.
"""
if not self.output_path.exists():
return {}
prompts: dict[str, str] = {}
with open(self.output_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if not line:
continue
try:
entry = json.loads(line)
if "filename" in entry and "prompt" in entry:
prompts[entry["filename"]] = entry["prompt"]
except json.JSONDecodeError:
pass
return prompts
def _save_entry(self, file_path: str, prompt: str) -> None:
"""Append one {filename, prompt} record to the JSONL output file."""
with open(self.output_path, "a", encoding="utf-8") as f:
f.write(
json.dumps({"filename": Path(file_path).stem, "prompt": prompt}) + "\n"
)
def run(self, port: int) -> None:
bottle.run(self.app, host="localhost", port=port, quiet=True)
# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main() -> None:
parser = argparse.ArgumentParser(
description="Browse images and assign text prompts via a web UI.",
)
parser.add_argument("images", nargs="+", help="Image paths to browse")
parser.add_argument(
"-o",
"--output-path",
default="prompts.jsonl",
metavar="FILE",
help="Output JSONL file path (default: prompts.jsonl)",
)
args = parser.parse_args()
images = [
str(Path(p).resolve())
for p in args.images
if Path(p).suffix.lower() in VALID_EXTENSIONS
]
if not images:
exts = ", ".join(sorted(VALID_EXTENSIONS))
print(
f"Error: no valid images found. Expected extensions: {exts}",
file=sys.stderr,
)
sys.exit(1)
print(f"Loaded {len(images)} image(s).")
port = find_free_port()
url = f"http://localhost:{port}"
print(f"Starting server at {url}")
threading.Thread(
target=lambda: (time.sleep(0.8), webbrowser.open(url)),
daemon=True,
).start()
PrompterApp(images, Path(args.output_path)).run(port)
if __name__ == "__main__":
main()