chore: Delete unused scripts
This commit is contained in:
1 parent
5ecff074af
commit
f449d692a8
34 files changed
-6047
No files matched your search
Whitespace-only changes.
Whitespace-only changes.
@@ -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()
|
|
||||||
@@ -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()
|
|
||||||
@@ -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
@@ -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
@@ -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()
|
|
||||||
@@ -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()
|
|
||||||
@@ -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())
|
|
||||||
@@ -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()
|
|
||||||
@@ -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()
|
|
||||||
@@ -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)
|
|
||||||
@@ -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()
|
|
||||||
@@ -1,5 +0,0 @@
|
|||||||
module github.com/abdusco/playground
|
|
||||||
|
|
||||||
go 1.23.0
|
|
||||||
|
|
||||||
require golang.org/x/sync v0.10.0
|
|
||||||
@@ -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
@@ -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
@@ -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
@@ -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
Binary file not shown.
-35
@@ -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())
|
|
||||||
@@ -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()
|
|
||||||
@@ -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()
|
|
||||||
@@ -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="Q7K9r45M9kh8bGqU8gIslAWIA1eaNo2M7w3xOZ6j",
|
|
||||||
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="Q7K9r45M9kh8bGqU8gIslAWIA1eaNo2M7w3xOZ6j",
|
|
||||||
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", "Q7K9r45M9kh8bGqU8gIslAWIA1eaNo2M7w3xOZ6j")
|
|
||||||
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)
|
|
||||||
@@ -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:")],
|
|
||||||
# )
|
|
||||||
# )
|
|
||||||
@@ -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"},
|
|
||||||
# )
|
|
||||||
# )
|
|
||||||
@@ -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",
|
|
||||||
# )
|
|
||||||
@@ -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()
|
|
||||||
@@ -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)
|
|
||||||
@@ -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])
|
|
||||||
@@ -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()
|
|
||||||
@@ -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
|
|
||||||
@@ -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]">✓ saved</span>
|
|
||||||
</div>
|
|
||||||
<textarea
|
|
||||||
x-ref="textarea"
|
|
||||||
x-model="customPrompt"
|
|
||||||
@focus="selectedOption = 'custom'"
|
|
||||||
placeholder="Enter a prompt…"
|
|
||||||
rows="5"
|
|
||||||
></textarea>
|
|
||||||
<button class="btn-save" @click.stop="saveAndNext()">Save & 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>
|
|
||||||
@@ -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()
|
|
||||||
Reference in new issue
Block a user