Import upstream v0.16.22, stripped
Upstream commit: 474dd0229cb20cf513036619781ed97bd8073c3f Enterprise-only files removed or emptied: 63 Enterprise-only snippets removed: 117 in 50 files Dangling module declarations removed: 5 Cargo edits turning enterprise off: 14 Verification: clean Enterprise feature gates left for rebuilt features: 19 in 18 files Produced by tools/fork/strip.py. The full report is in docs/fork/strip-reports/ on main.
This commit is contained in:
@@ -0,0 +1,55 @@
|
||||
import imaplib
|
||||
import socket
|
||||
import time
|
||||
import threading
|
||||
from email.message import Message
|
||||
from email.utils import formatdate
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
def append_message(thread_id, start, end):
|
||||
conn = imaplib.IMAP4('localhost')
|
||||
conn.login('john', '12345')
|
||||
conn.socket().setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
start_time = time.time()
|
||||
|
||||
base_date = datetime(2000, 1, 1)
|
||||
|
||||
for n in range(start, end):
|
||||
current_date = base_date + timedelta(hours=n)
|
||||
|
||||
msg = Message()
|
||||
msg['From'] = '[email protected]'
|
||||
msg['To'] = '[email protected]'
|
||||
msg['Message-Id'] = f'unique.message.id.{n}@nowhere'
|
||||
msg['Date'] = formatdate(time.mktime(current_date.timetuple()), localtime=False, usegmt=True)
|
||||
msg['Subject'] = f"This is message #{n}"
|
||||
msg.set_payload('...nothing...')
|
||||
|
||||
response_code, response_details = conn.append('INBOX', '', imaplib.Time2Internaldate(time.mktime(current_date.timetuple())), str(msg).encode('utf-8'))
|
||||
if response_code != 'OK':
|
||||
print(f'Thread {thread_id}: Error while appending message #{n}: {response_code} {response_details}')
|
||||
break
|
||||
if n != 0 and n % 100 == 0:
|
||||
elapsed_time = (time.time() - start_time) * 1000
|
||||
print(f'Thread {thread_id}: Inserting batch {n} took {elapsed_time} ms.', flush=True)
|
||||
start_time = time.time()
|
||||
|
||||
conn.logout()
|
||||
|
||||
num_threads = 5
|
||||
num_messages = 10000
|
||||
messages_per_thread = num_messages // num_threads
|
||||
|
||||
threads = []
|
||||
for i in range(num_threads):
|
||||
start = i * messages_per_thread
|
||||
end = start + messages_per_thread
|
||||
thread = threading.Thread(target=append_message, args=(i, start, end))
|
||||
threads.append(thread)
|
||||
thread.start()
|
||||
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
|
||||
print("All messages appended.")
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
import imaplib
|
||||
import socket
|
||||
import time
|
||||
from email.message import Message
|
||||
from email.utils import formatdate
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
conn = imaplib.IMAP4('localhost')
|
||||
conn.login('john', '12345')
|
||||
conn.socket().setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
|
||||
current_date = datetime.now()
|
||||
timestamp = current_date.timestamp()
|
||||
|
||||
msg = Message()
|
||||
msg['From'] = '[email protected]'
|
||||
msg['To'] = '[email protected]'
|
||||
msg['Message-Id'] = f'unique.message.id.{current_date}@nowhere'
|
||||
msg['Date'] = formatdate(time.mktime(current_date.timetuple()), localtime=False, usegmt=True)
|
||||
msg['Subject'] = f"This is message #{timestamp}"
|
||||
msg.set_payload('...nothing...')
|
||||
|
||||
response_code, response_details = conn.append('INBOX', '', imaplib.Time2Internaldate(time.mktime(current_date.timetuple())), str(msg).encode('utf-8'))
|
||||
if response_code != 'OK':
|
||||
print(f'Error while appending message: {response_code} {response_details}')
|
||||
|
||||
print("Message appended.")
|
||||
conn.logout()
|
||||
@@ -0,0 +1,627 @@
|
||||
import argparse
|
||||
import math
|
||||
import multiprocessing
|
||||
import os
|
||||
import queue
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import smtplib
|
||||
import ssl
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from email.utils import formatdate, make_msgid
|
||||
|
||||
DEFAULT_HOST = "127.0.0.1"
|
||||
DEFAULT_PORT = 25
|
||||
DEFAULT_THREADS = 5
|
||||
DEFAULT_PROCESSES = 1
|
||||
DEFAULT_MESSAGES = 100
|
||||
DEFAULT_MIN_SIZE = 1024
|
||||
DEFAULT_MAX_SIZE = 51200
|
||||
DEFAULT_POOL_SIZE = 64
|
||||
DEFAULT_SENDER = "[email protected]"
|
||||
DEFAULT_USERS_FILE = "users.txt"
|
||||
DEFAULT_DICTIONARY = "/usr/share/dict/words"
|
||||
DEFAULT_TIMEOUT = 60
|
||||
LINE_WIDTH = 72
|
||||
|
||||
FALLBACK_WORDS = (
|
||||
"lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod "
|
||||
"tempor incididunt ut labore et dolore magna aliqua enim ad minim veniam "
|
||||
"quis nostrud exercitation ullamco laboris nisi aliquip ex ea commodo"
|
||||
).split()
|
||||
|
||||
WORDS = FALLBACK_WORDS
|
||||
|
||||
DOT_LINE = re.compile(br"(?m)^\.")
|
||||
STOP_EVENT = threading.Event()
|
||||
|
||||
|
||||
class SmtpError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class AsyncLogger:
|
||||
def __init__(self, enabled):
|
||||
self.enabled = enabled
|
||||
self._queue = queue.Queue() if enabled else None
|
||||
self._thread = None
|
||||
|
||||
def start(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
self._thread = threading.Thread(target=self._run, daemon=True)
|
||||
self._thread.start()
|
||||
|
||||
def log(self, message):
|
||||
if self.enabled:
|
||||
self._queue.put(message)
|
||||
|
||||
def _run(self):
|
||||
while True:
|
||||
message = self._queue.get()
|
||||
if message is None:
|
||||
break
|
||||
print(message, file=sys.stderr, flush=True)
|
||||
|
||||
def stop(self):
|
||||
if not self.enabled:
|
||||
return
|
||||
self._queue.put(None)
|
||||
if self._thread is not None:
|
||||
self._thread.join()
|
||||
|
||||
|
||||
def load_words(path):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="ignore") as file:
|
||||
words = [w.strip() for w in file if w.strip().isalpha()]
|
||||
except OSError:
|
||||
words = []
|
||||
if not words:
|
||||
print(
|
||||
f"WARNING: word list '{path}' not found or empty; "
|
||||
f"falling back to built-in lorem ipsum words. "
|
||||
f"Override with --dict <path>.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return list(FALLBACK_WORDS)
|
||||
return words
|
||||
|
||||
|
||||
def random_subject():
|
||||
return " ".join(random.choices(WORDS, k=random.randint(3, 10)))
|
||||
|
||||
|
||||
def random_body(target_size):
|
||||
lines = []
|
||||
total = 0
|
||||
line = ""
|
||||
while total < target_size:
|
||||
word = random.choice(WORDS)
|
||||
if line and len(line) + 1 + len(word) > LINE_WIDTH:
|
||||
lines.append(line)
|
||||
total += len(line) + 2
|
||||
line = word
|
||||
elif line:
|
||||
line = f"{line} {word}"
|
||||
else:
|
||||
line = word
|
||||
if line:
|
||||
lines.append(line)
|
||||
return "\r\n".join(lines) + "\r\n"
|
||||
|
||||
|
||||
def quote_periods(data):
|
||||
return DOT_LINE.sub(b"..", data)
|
||||
|
||||
|
||||
def build_headers(sender, recipient):
|
||||
return (
|
||||
f"From: {sender}\r\n"
|
||||
f"To: {recipient}\r\n"
|
||||
f"Subject: {random_subject()}\r\n"
|
||||
f"Date: {formatdate(localtime=True)}\r\n"
|
||||
f"Message-ID: {make_msgid(domain='stress.test')}\r\n"
|
||||
f"MIME-Version: 1.0\r\n"
|
||||
f"Content-Type: text/plain; charset=us-ascii\r\n"
|
||||
f"\r\n"
|
||||
).encode("ascii", "replace")
|
||||
|
||||
|
||||
def build_body(size):
|
||||
body = quote_periods(random_body(size).encode("ascii", "replace"))
|
||||
if not body.endswith(b"\r\n"):
|
||||
body += b"\r\n"
|
||||
return body
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
backend = "memory"
|
||||
|
||||
def __init__(self):
|
||||
self._items = []
|
||||
|
||||
def add(self, data):
|
||||
self._items.append(data)
|
||||
|
||||
def get(self, index):
|
||||
return self._items[index]
|
||||
|
||||
def __len__(self):
|
||||
return len(self._items)
|
||||
|
||||
def cleanup(self):
|
||||
self._items = []
|
||||
|
||||
|
||||
class DiskStore:
|
||||
def __init__(self, root):
|
||||
self._dir = tempfile.mkdtemp(prefix="smtp_stress_", dir=root)
|
||||
self.backend = self._dir
|
||||
self._paths = []
|
||||
|
||||
def add(self, data):
|
||||
path = os.path.join(self._dir, f"msg_{len(self._paths):09d}.eml")
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(data)
|
||||
self._paths.append(path)
|
||||
|
||||
def get(self, index):
|
||||
with open(self._paths[index], "rb") as handle:
|
||||
return handle.read()
|
||||
|
||||
def __len__(self):
|
||||
return len(self._paths)
|
||||
|
||||
def cleanup(self):
|
||||
shutil.rmtree(self._dir, ignore_errors=True)
|
||||
|
||||
|
||||
def build_body_store(ctx):
|
||||
count = min(ctx.pool_size, ctx.messages)
|
||||
if ctx.spool_dir is not None:
|
||||
store = DiskStore(ctx.spool_dir)
|
||||
else:
|
||||
store = MemoryStore()
|
||||
for _ in range(count):
|
||||
if ctx.fixed_size is not None:
|
||||
size = ctx.fixed_size
|
||||
else:
|
||||
size = random.randint(ctx.min_size, ctx.max_size)
|
||||
store.add(build_body(size))
|
||||
return store
|
||||
|
||||
|
||||
def make_tls_context():
|
||||
context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
return context
|
||||
|
||||
|
||||
class Stats:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.latencies = []
|
||||
self.errors = 0
|
||||
self.bytes = 0
|
||||
|
||||
def record(self, latency_ms, num_bytes):
|
||||
with self._lock:
|
||||
self.latencies.append(latency_ms)
|
||||
self.bytes += num_bytes
|
||||
|
||||
def record_error(self):
|
||||
with self._lock:
|
||||
self.errors += 1
|
||||
|
||||
def snapshot(self):
|
||||
with self._lock:
|
||||
return list(self.latencies), self.errors, self.bytes
|
||||
|
||||
|
||||
class Counter:
|
||||
def __init__(self, total):
|
||||
self._lock = threading.Lock()
|
||||
self._remaining = total
|
||||
|
||||
def claim(self):
|
||||
with self._lock:
|
||||
if self._remaining <= 0:
|
||||
return False
|
||||
self._remaining -= 1
|
||||
return True
|
||||
|
||||
|
||||
def read_recipients(file_path):
|
||||
recipients = []
|
||||
try:
|
||||
with open(file_path, "r") as file:
|
||||
for line in file:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
recipients.append(line.split(":", 1)[0])
|
||||
except OSError as e:
|
||||
raise SystemExit(f"Could not read recipients from '{file_path}': {e}")
|
||||
if not recipients:
|
||||
raise SystemExit(f"No recipients found in '{file_path}'.")
|
||||
return recipients
|
||||
|
||||
|
||||
def connect(ctx):
|
||||
server = smtplib.SMTP(ctx.host, ctx.port, timeout=ctx.timeout)
|
||||
server.ehlo()
|
||||
if ctx.starttls:
|
||||
if not server.has_extn("starttls"):
|
||||
server.quit()
|
||||
raise SmtpError("server does not advertise STARTTLS")
|
||||
server.starttls(context=ctx.tls_context)
|
||||
server.ehlo()
|
||||
return server
|
||||
|
||||
|
||||
def send_one(server, sender, recipient, header, body):
|
||||
code, resp = server.mail(sender)
|
||||
if code != 250:
|
||||
server.rset()
|
||||
raise SmtpError(f"MAIL FROM rejected: {code} {resp!r}")
|
||||
code, resp = server.rcpt(recipient)
|
||||
if code not in (250, 251):
|
||||
server.rset()
|
||||
raise SmtpError(f"RCPT TO rejected: {code} {resp!r}")
|
||||
code, resp = server.docmd("DATA")
|
||||
if code != 354:
|
||||
raise SmtpError(f"DATA rejected: {code} {resp!r}")
|
||||
server.send(header)
|
||||
server.send(body)
|
||||
start = time.monotonic()
|
||||
server.send(b".\r\n")
|
||||
code, resp = server.getreply()
|
||||
elapsed_ms = (time.monotonic() - start) * 1000
|
||||
if code != 250:
|
||||
raise SmtpError(f"message rejected: {code} {resp!r}")
|
||||
return elapsed_ms
|
||||
|
||||
|
||||
def worker(ctx, counter, recipients, stats, store, logger):
|
||||
pool_len = len(store)
|
||||
while not STOP_EVENT.is_set() and counter.claim():
|
||||
server = None
|
||||
try:
|
||||
server = connect(ctx)
|
||||
recipient = random.choice(recipients)
|
||||
body = store.get(random.randrange(pool_len))
|
||||
header = build_headers(ctx.sender, recipient)
|
||||
elapsed_ms = send_one(server, ctx.sender, recipient, header, body)
|
||||
num_bytes = len(header) + len(body)
|
||||
stats.record(elapsed_ms, num_bytes)
|
||||
if logger.enabled:
|
||||
logger.log(f"OK {elapsed_ms:9.2f}ms {num_bytes:>9}B -> {recipient}")
|
||||
except (SmtpError, smtplib.SMTPException, OSError) as e:
|
||||
stats.record_error()
|
||||
if logger.enabled:
|
||||
logger.log(f"ERR {e}")
|
||||
finally:
|
||||
if server is not None:
|
||||
try:
|
||||
server.quit()
|
||||
except Exception:
|
||||
try:
|
||||
server.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def percentile(sorted_values, pct):
|
||||
if not sorted_values:
|
||||
return 0.0
|
||||
if len(sorted_values) == 1:
|
||||
return sorted_values[0]
|
||||
rank = (len(sorted_values) - 1) * (pct / 100.0)
|
||||
low = math.floor(rank)
|
||||
high = math.ceil(rank)
|
||||
if low == high:
|
||||
return sorted_values[int(rank)]
|
||||
return sorted_values[low] * (high - rank) + sorted_values[high] * (rank - low)
|
||||
|
||||
|
||||
def stddev(values, mean):
|
||||
if len(values) < 2:
|
||||
return 0.0
|
||||
variance = sum((v - mean) ** 2 for v in values) / (len(values) - 1)
|
||||
return math.sqrt(variance)
|
||||
|
||||
|
||||
def print_report(
|
||||
latencies,
|
||||
errors,
|
||||
total_bytes,
|
||||
send_seconds,
|
||||
gen_seconds,
|
||||
pool_count,
|
||||
workers,
|
||||
storage,
|
||||
report_header=None,
|
||||
):
|
||||
count = len(latencies)
|
||||
mb = total_bytes / (1024 * 1024)
|
||||
throughput = count / send_seconds if send_seconds > 0 else 0.0
|
||||
mb_per_sec = mb / send_seconds if send_seconds > 0 else 0.0
|
||||
|
||||
line = "-" * 60
|
||||
print()
|
||||
if report_header:
|
||||
print(report_header)
|
||||
print(line)
|
||||
print("SMTP ingestion stress test report")
|
||||
print(line)
|
||||
print(f"{'Workers':<26}{workers}")
|
||||
print(f"{'Message store':<26}{storage}")
|
||||
print(f"{'Messages OK':<26}{count}")
|
||||
print(f"{'Messages failed':<26}{errors}")
|
||||
print(f"{'Bodies pregenerated':<26}{pool_count}")
|
||||
print(f"{'Pool gen time (s)':<26}{gen_seconds:.2f}")
|
||||
print(f"{'Send wall time (s)':<26}{send_seconds:.2f}")
|
||||
print(f"{'Throughput (msg/s)':<26}{throughput:.2f}")
|
||||
print(f"{'Data sent (MB)':<26}{mb:.2f}")
|
||||
print(f"{'Data rate (MB/s)':<26}{mb_per_sec:.2f}")
|
||||
print(line)
|
||||
print("Ingestion time (DATA terminator to server OK), milliseconds")
|
||||
print(line)
|
||||
if count:
|
||||
ordered = sorted(latencies)
|
||||
mean = sum(ordered) / count
|
||||
rows = [
|
||||
("min", ordered[0]),
|
||||
("max", ordered[-1]),
|
||||
("avg", mean),
|
||||
("median", percentile(ordered, 50)),
|
||||
("p95", percentile(ordered, 95)),
|
||||
("p99", percentile(ordered, 99)),
|
||||
("stddev", stddev(ordered, mean)),
|
||||
]
|
||||
for name, value in rows:
|
||||
print(f"{name:<22}{value:.2f}")
|
||||
else:
|
||||
print("no messages were ingested")
|
||||
print(line)
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
class Context:
|
||||
def __init__(self, args, messages):
|
||||
self.host = args.host
|
||||
self.port = args.port
|
||||
self.threads = args.threads
|
||||
self.sender = args.sender
|
||||
self.starttls = not args.no_starttls
|
||||
self.timeout = args.timeout
|
||||
self.min_size = args.min_size
|
||||
self.max_size = args.max_size
|
||||
self.fixed_size = args.size
|
||||
self.pool_size = args.pool_size
|
||||
self.spool_dir = args.spool_dir
|
||||
self.messages = messages
|
||||
self.tls_context = make_tls_context() if self.starttls else None
|
||||
|
||||
|
||||
def run_threads(ctx, recipients, message_count, store, logger):
|
||||
stats = Stats()
|
||||
counter = Counter(message_count)
|
||||
threads = [
|
||||
threading.Thread(
|
||||
target=worker,
|
||||
args=(ctx, counter, recipients, stats, store, logger),
|
||||
daemon=True,
|
||||
)
|
||||
for _ in range(ctx.threads)
|
||||
]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
try:
|
||||
while any(t.is_alive() for t in threads):
|
||||
for t in threads:
|
||||
t.join(timeout=0.2)
|
||||
except KeyboardInterrupt:
|
||||
logger.log("Stopping...")
|
||||
STOP_EVENT.set()
|
||||
for t in threads:
|
||||
t.join()
|
||||
return stats
|
||||
|
||||
|
||||
def child_main(args, recipients, message_count, barrier, result_queue):
|
||||
global WORDS
|
||||
WORDS = load_words(args.dict)
|
||||
ctx = Context(args, message_count)
|
||||
logger = AsyncLogger(not args.quiet)
|
||||
store = build_body_store(ctx)
|
||||
try:
|
||||
logger.start()
|
||||
try:
|
||||
barrier.wait()
|
||||
except threading.BrokenBarrierError:
|
||||
result_queue.put(([], 0, 0, len(store)))
|
||||
return
|
||||
stats = run_threads(ctx, recipients, message_count, store, logger)
|
||||
logger.stop()
|
||||
latencies, errors, total_bytes = stats.snapshot()
|
||||
result_queue.put((latencies, errors, total_bytes, len(store)))
|
||||
finally:
|
||||
store.cleanup()
|
||||
|
||||
|
||||
def distribute(total, parts):
|
||||
base, remainder = divmod(total, parts)
|
||||
return [base + (1 if i < remainder else 0) for i in range(parts)]
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Concurrent SMTP ingestion stress test over port 25 with STARTTLS."
|
||||
)
|
||||
parser.add_argument("--host", default=DEFAULT_HOST)
|
||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
||||
parser.add_argument("--threads", type=int, default=DEFAULT_THREADS)
|
||||
parser.add_argument(
|
||||
"--processes",
|
||||
type=int,
|
||||
default=DEFAULT_PROCESSES,
|
||||
help="Worker processes to spawn (each runs --threads threads). Scales past the GIL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--messages",
|
||||
type=int,
|
||||
default=DEFAULT_MESSAGES,
|
||||
help="Total messages to send, distributed across the threads.",
|
||||
)
|
||||
parser.add_argument("--sender", default=DEFAULT_SENDER, help="Envelope MAIL FROM address.")
|
||||
parser.add_argument("--users-file", default=DEFAULT_USERS_FILE)
|
||||
parser.add_argument(
|
||||
"--size",
|
||||
type=int,
|
||||
help="Fixed message body size in bytes; overrides --min-size/--max-size.",
|
||||
)
|
||||
parser.add_argument("--min-size", type=int, default=DEFAULT_MIN_SIZE)
|
||||
parser.add_argument("--max-size", type=int, default=DEFAULT_MAX_SIZE)
|
||||
parser.add_argument(
|
||||
"--pool-size",
|
||||
type=int,
|
||||
default=DEFAULT_POOL_SIZE,
|
||||
help="Distinct message bodies pregenerated before timing (reused at random). "
|
||||
"Each sent message gets a fresh unique Message-ID regardless of this.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--spool-dir",
|
||||
nargs="?",
|
||||
const=tempfile.gettempdir(),
|
||||
default=None,
|
||||
help="Spool pregenerated messages to disk instead of memory. "
|
||||
"With no value uses the system temp dir; pass a path to override.",
|
||||
)
|
||||
parser.add_argument("--dict", default=DEFAULT_DICTIONARY)
|
||||
parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT)
|
||||
parser.add_argument(
|
||||
"--header",
|
||||
help="Optional header line printed at the top of the final report on stdout.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-starttls",
|
||||
action="store_true",
|
||||
help="Send over plaintext instead of upgrading with STARTTLS.",
|
||||
)
|
||||
parser.add_argument("--quiet", action="store_true")
|
||||
args = parser.parse_args()
|
||||
if args.threads < 1:
|
||||
parser.error("--threads must be at least 1")
|
||||
if args.processes < 1:
|
||||
parser.error("--processes must be at least 1")
|
||||
if args.messages < 1:
|
||||
parser.error("--messages must be at least 1")
|
||||
if args.pool_size < 1:
|
||||
parser.error("--pool-size must be at least 1")
|
||||
if args.size is None and args.min_size > args.max_size:
|
||||
parser.error("--min-size must not exceed --max-size")
|
||||
return args
|
||||
|
||||
|
||||
def run_single_process(args, recipients):
|
||||
global WORDS
|
||||
WORDS = load_words(args.dict)
|
||||
ctx = Context(args, args.messages)
|
||||
logger = AsyncLogger(not args.quiet)
|
||||
|
||||
gen_start = time.monotonic()
|
||||
store = build_body_store(ctx)
|
||||
gen_seconds = time.monotonic() - gen_start
|
||||
|
||||
try:
|
||||
logger.start()
|
||||
start = time.monotonic()
|
||||
stats = run_threads(ctx, recipients, args.messages, store, logger)
|
||||
send_seconds = time.monotonic() - start
|
||||
logger.stop()
|
||||
latencies, errors, total_bytes = stats.snapshot()
|
||||
workers = f"1 process x {args.threads} threads"
|
||||
print_report(
|
||||
latencies, errors, total_bytes, send_seconds, gen_seconds, len(store),
|
||||
workers, store.backend, args.header,
|
||||
)
|
||||
finally:
|
||||
store.cleanup()
|
||||
|
||||
|
||||
def run_multi_process(args, recipients):
|
||||
nproc = min(args.processes, args.messages)
|
||||
shares = distribute(args.messages, nproc)
|
||||
barrier = multiprocessing.Barrier(nproc + 1)
|
||||
result_queue = multiprocessing.Queue()
|
||||
procs = []
|
||||
for share in shares:
|
||||
proc = multiprocessing.Process(
|
||||
target=child_main,
|
||||
args=(args, recipients, share, barrier, result_queue),
|
||||
daemon=False,
|
||||
)
|
||||
proc.start()
|
||||
procs.append(proc)
|
||||
|
||||
gen_start = time.monotonic()
|
||||
interrupted = False
|
||||
try:
|
||||
barrier.wait()
|
||||
except KeyboardInterrupt:
|
||||
interrupted = True
|
||||
barrier.abort()
|
||||
gen_seconds = time.monotonic() - gen_start
|
||||
|
||||
start = time.monotonic()
|
||||
results = []
|
||||
try:
|
||||
for _ in procs:
|
||||
results.append(result_queue.get())
|
||||
except KeyboardInterrupt:
|
||||
interrupted = True
|
||||
for proc in procs:
|
||||
proc.terminate()
|
||||
send_seconds = time.monotonic() - start
|
||||
|
||||
for proc in procs:
|
||||
proc.join()
|
||||
|
||||
latencies = []
|
||||
errors = 0
|
||||
total_bytes = 0
|
||||
pool_count = 0
|
||||
for lat, err, nbytes, pool_len in results:
|
||||
latencies.extend(lat)
|
||||
errors += err
|
||||
total_bytes += nbytes
|
||||
pool_count += pool_len
|
||||
|
||||
if interrupted:
|
||||
print("Interrupted.", file=sys.stderr, flush=True)
|
||||
workers = f"{nproc} processes x {args.threads} threads"
|
||||
storage = "memory" if args.spool_dir is None else f"disk ({args.spool_dir})"
|
||||
print_report(
|
||||
latencies, errors, total_bytes, send_seconds, gen_seconds, pool_count,
|
||||
workers, storage, args.header,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
recipients = read_recipients(args.users_file)
|
||||
if args.processes == 1:
|
||||
run_single_process(args, recipients)
|
||||
else:
|
||||
run_multi_process(args, recipients)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,490 @@
|
||||
import argparse
|
||||
import imaplib
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import smtplib
|
||||
import ssl
|
||||
import threading
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from email.mime.text import MIMEText
|
||||
|
||||
DEFAULT_SMTP_SERVER = "127.0.0.1"
|
||||
DEFAULT_SMTP_PORT = 465
|
||||
DEFAULT_IMAP_SERVER = "127.0.0.1"
|
||||
DEFAULT_IMAP_PORT = 993
|
||||
DEFAULT_THREADS = 5
|
||||
DEFAULT_RUNS = 10
|
||||
DEFAULT_DICTIONARY = "/usr/share/dict/words"
|
||||
|
||||
FALLBACK_WORDS = (
|
||||
"lorem ipsum dolor sit amet consectetur adipiscing elit sed do eiusmod "
|
||||
"tempor incididunt ut labore et dolore magna aliqua enim ad minim veniam "
|
||||
"quis nostrud exercitation ullamco laboris nisi aliquip ex ea commodo"
|
||||
).split()
|
||||
|
||||
WORDS = FALLBACK_WORDS
|
||||
|
||||
SMTP_SEND = "SMTP SEND"
|
||||
IMAP_APPEND = "IMAP APPEND"
|
||||
IMAP_FETCH = "IMAP FETCH"
|
||||
IMAP_DELETE = "IMAP DELETE"
|
||||
ACTIONS = (SMTP_SEND, IMAP_APPEND, IMAP_FETCH, IMAP_DELETE)
|
||||
|
||||
|
||||
class Stats:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._latencies = defaultdict(list)
|
||||
self._errors = defaultdict(int)
|
||||
self._skips = defaultdict(int)
|
||||
self._bytes = defaultdict(int)
|
||||
|
||||
def record(self, action, latency_ms, num_bytes=0):
|
||||
with self._lock:
|
||||
self._latencies[action].append(latency_ms)
|
||||
self._bytes[action] += num_bytes
|
||||
|
||||
def record_error(self, action):
|
||||
with self._lock:
|
||||
self._errors[action] += 1
|
||||
|
||||
def record_skip(self, action):
|
||||
with self._lock:
|
||||
self._skips[action] += 1
|
||||
|
||||
def snapshot(self):
|
||||
with self._lock:
|
||||
return (
|
||||
{k: list(v) for k, v in self._latencies.items()},
|
||||
dict(self._errors),
|
||||
dict(self._skips),
|
||||
dict(self._bytes),
|
||||
)
|
||||
|
||||
|
||||
PRINT_LOCK = threading.Lock()
|
||||
STOP_EVENT = threading.Event()
|
||||
|
||||
|
||||
def read_credentials(file_path):
|
||||
if not os.path.exists(file_path):
|
||||
raise SystemExit(
|
||||
f"Credentials file '{file_path}' not found. "
|
||||
f"Run stress_test_prepare.py first to create users."
|
||||
)
|
||||
credentials = []
|
||||
with open(file_path, "r") as file:
|
||||
for line in file:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
parts = line.split(":", 1)
|
||||
if len(parts) != 2:
|
||||
continue
|
||||
credentials.append((parts[0], parts[1]))
|
||||
if not credentials:
|
||||
raise SystemExit(f"No valid credentials found in '{file_path}'.")
|
||||
return credentials
|
||||
|
||||
|
||||
def allow_invalid_certificates():
|
||||
context = ssl.create_default_context()
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
return context
|
||||
|
||||
|
||||
def load_words(path):
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="ignore") as file:
|
||||
words = [w.strip() for w in file if w.strip().isalpha()]
|
||||
except OSError:
|
||||
words = []
|
||||
if not words:
|
||||
print(
|
||||
f"WARNING: word list '{path}' not found or empty; "
|
||||
f"falling back to built-in lorem ipsum words. "
|
||||
f"Override with --dict <path>."
|
||||
)
|
||||
return list(FALLBACK_WORDS)
|
||||
return words
|
||||
|
||||
|
||||
def random_words(min_size, max_size):
|
||||
target = random.randint(min_size, max_size)
|
||||
parts = []
|
||||
length = 0
|
||||
while length < target:
|
||||
for word in random.choices(WORDS, k=64):
|
||||
parts.append(word)
|
||||
length += len(word) + 1
|
||||
if length >= target:
|
||||
break
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
def generate_email(username, recipient, max_content_size):
|
||||
subject = random_words(10, 100)
|
||||
content = random_words(100, max_content_size)
|
||||
message = MIMEText(content)
|
||||
message["Subject"] = subject
|
||||
message["From"] = username
|
||||
message["To"] = recipient
|
||||
return message.as_string()
|
||||
|
||||
|
||||
def log_ok(stats, action, latency_ms, detail="", num_bytes=0, verbose=True):
|
||||
stats.record(action, latency_ms, num_bytes)
|
||||
if verbose:
|
||||
with PRINT_LOCK:
|
||||
print(f"OK {latency_ms:9.2f}ms {action} {detail}")
|
||||
|
||||
|
||||
def log_err(stats, action, error, verbose=True):
|
||||
stats.record_error(action)
|
||||
if verbose:
|
||||
with PRINT_LOCK:
|
||||
print(f"ERR {action} {error}")
|
||||
|
||||
|
||||
def smtp_send_message(ctx, username, password, recipient):
|
||||
try:
|
||||
with smtplib.SMTP_SSL(
|
||||
ctx.smtp_server, ctx.smtp_port, context=allow_invalid_certificates()
|
||||
) as server:
|
||||
server.login(username, password)
|
||||
payload = generate_email(username, recipient, ctx.max_content_size)
|
||||
start_time = time.monotonic()
|
||||
server.sendmail(username, recipient, payload)
|
||||
elapsed_ms = (time.monotonic() - start_time) * 1000
|
||||
log_ok(
|
||||
ctx.stats,
|
||||
SMTP_SEND,
|
||||
elapsed_ms,
|
||||
f"{username} -> {recipient}",
|
||||
len(payload),
|
||||
ctx.verbose,
|
||||
)
|
||||
except Exception as e:
|
||||
log_err(ctx.stats, SMTP_SEND, e, ctx.verbose)
|
||||
|
||||
|
||||
def imap_append_message(ctx, username, password, recipient):
|
||||
try:
|
||||
with imaplib.IMAP4_SSL(
|
||||
ctx.imap_server, ctx.imap_port, ssl_context=allow_invalid_certificates()
|
||||
) as imap:
|
||||
imap.login(username, password)
|
||||
payload = generate_email(username, recipient, ctx.max_content_size).encode("utf-8")
|
||||
start_time = time.monotonic()
|
||||
imap.append("INBOX", None, imaplib.Time2Internaldate(time.time()), payload)
|
||||
elapsed_ms = (time.monotonic() - start_time) * 1000
|
||||
log_ok(ctx.stats, IMAP_APPEND, elapsed_ms, username, len(payload), ctx.verbose)
|
||||
except Exception as e:
|
||||
log_err(ctx.stats, IMAP_APPEND, e, ctx.verbose)
|
||||
|
||||
|
||||
def imap_list_fetch(ctx, username, password, recipient):
|
||||
try:
|
||||
with imaplib.IMAP4_SSL(
|
||||
ctx.imap_server, ctx.imap_port, ssl_context=allow_invalid_certificates()
|
||||
) as imap:
|
||||
imap.login(username, password)
|
||||
imap.select("INBOX")
|
||||
start_time = time.monotonic()
|
||||
typ, data = imap.search(None, "ALL")
|
||||
if data and data[0]:
|
||||
messages = data[0].split()
|
||||
random_msg_num = random.choice(messages)
|
||||
imap.fetch(random_msg_num, "(RFC822)")
|
||||
elapsed_ms = (time.monotonic() - start_time) * 1000
|
||||
log_ok(
|
||||
ctx.stats,
|
||||
IMAP_FETCH,
|
||||
elapsed_ms,
|
||||
f"{username} {random_msg_num.decode()}",
|
||||
verbose=ctx.verbose,
|
||||
)
|
||||
else:
|
||||
ctx.stats.record_skip(IMAP_FETCH)
|
||||
except Exception as e:
|
||||
log_err(ctx.stats, IMAP_FETCH, e, ctx.verbose)
|
||||
|
||||
|
||||
def imap_delete_message(ctx, username, password, recipient):
|
||||
try:
|
||||
with imaplib.IMAP4_SSL(
|
||||
ctx.imap_server, ctx.imap_port, ssl_context=allow_invalid_certificates()
|
||||
) as imap:
|
||||
imap.login(username, password)
|
||||
imap.select("INBOX")
|
||||
start_time = time.monotonic()
|
||||
typ, data = imap.search(None, "ALL")
|
||||
if data and data[0]:
|
||||
messages = data[0].split()
|
||||
random_msg_num = random.choice(messages)
|
||||
imap.store(random_msg_num, "+FLAGS", "\\Deleted")
|
||||
imap.expunge()
|
||||
elapsed_ms = (time.monotonic() - start_time) * 1000
|
||||
log_ok(
|
||||
ctx.stats,
|
||||
IMAP_DELETE,
|
||||
elapsed_ms,
|
||||
f"{username} {random_msg_num.decode()}",
|
||||
verbose=ctx.verbose,
|
||||
)
|
||||
else:
|
||||
ctx.stats.record_skip(IMAP_DELETE)
|
||||
except Exception as e:
|
||||
log_err(ctx.stats, IMAP_DELETE, e, ctx.verbose)
|
||||
|
||||
|
||||
ACTION_FUNCS = (
|
||||
smtp_send_message,
|
||||
imap_append_message,
|
||||
imap_list_fetch,
|
||||
imap_delete_message,
|
||||
)
|
||||
|
||||
|
||||
def pick_recipient(credentials, sender):
|
||||
if len(credentials) == 1:
|
||||
return credentials[0][0]
|
||||
while True:
|
||||
recipient = random.choice(credentials)[0]
|
||||
if recipient != sender:
|
||||
return recipient
|
||||
|
||||
|
||||
def perform_random_action(ctx, credentials):
|
||||
username, password = random.choice(credentials)
|
||||
recipient = pick_recipient(credentials, username)
|
||||
action = random.choice(ACTION_FUNCS)
|
||||
action(ctx, username, password, recipient)
|
||||
|
||||
|
||||
def thread_function(ctx, credentials):
|
||||
count = 0
|
||||
while not STOP_EVENT.is_set():
|
||||
if ctx.runs is not None and count >= ctx.runs:
|
||||
break
|
||||
perform_random_action(ctx, credentials)
|
||||
count += 1
|
||||
|
||||
|
||||
def percentile(sorted_values, pct):
|
||||
if not sorted_values:
|
||||
return 0.0
|
||||
if len(sorted_values) == 1:
|
||||
return sorted_values[0]
|
||||
rank = (len(sorted_values) - 1) * (pct / 100.0)
|
||||
low = math.floor(rank)
|
||||
high = math.ceil(rank)
|
||||
if low == high:
|
||||
return sorted_values[int(rank)]
|
||||
return sorted_values[low] * (high - rank) + sorted_values[high] * (rank - low)
|
||||
|
||||
|
||||
def stddev(values, mean):
|
||||
if len(values) < 2:
|
||||
return 0.0
|
||||
variance = sum((v - mean) ** 2 for v in values) / (len(values) - 1)
|
||||
return math.sqrt(variance)
|
||||
|
||||
|
||||
def summarize(action, latencies, errors, skips, total_bytes, wall_seconds):
|
||||
count = len(latencies)
|
||||
summary = {
|
||||
"action": action,
|
||||
"count": count,
|
||||
"errors": errors,
|
||||
"skips": skips,
|
||||
"mb": total_bytes / (1024 * 1024),
|
||||
}
|
||||
if count == 0:
|
||||
for key in ("min", "max", "avg", "median", "p95", "p99", "stddev", "ops"):
|
||||
summary[key] = 0.0
|
||||
return summary
|
||||
ordered = sorted(latencies)
|
||||
mean = sum(ordered) / count
|
||||
summary.update(
|
||||
{
|
||||
"min": ordered[0],
|
||||
"max": ordered[-1],
|
||||
"avg": mean,
|
||||
"median": percentile(ordered, 50),
|
||||
"p95": percentile(ordered, 95),
|
||||
"p99": percentile(ordered, 99),
|
||||
"stddev": stddev(ordered, mean),
|
||||
"ops": count / wall_seconds if wall_seconds > 0 else 0.0,
|
||||
}
|
||||
)
|
||||
return summary
|
||||
|
||||
|
||||
def print_report(stats, wall_seconds):
|
||||
latencies, errors, skips, byte_counts = stats.snapshot()
|
||||
|
||||
rows = []
|
||||
all_latencies = []
|
||||
total_errors = 0
|
||||
total_skips = 0
|
||||
total_bytes = 0
|
||||
for action in ACTIONS:
|
||||
action_latencies = latencies.get(action, [])
|
||||
all_latencies.extend(action_latencies)
|
||||
total_errors += errors.get(action, 0)
|
||||
total_skips += skips.get(action, 0)
|
||||
total_bytes += byte_counts.get(action, 0)
|
||||
rows.append(
|
||||
summarize(
|
||||
action,
|
||||
action_latencies,
|
||||
errors.get(action, 0),
|
||||
skips.get(action, 0),
|
||||
byte_counts.get(action, 0),
|
||||
wall_seconds,
|
||||
)
|
||||
)
|
||||
rows.append(
|
||||
summarize(
|
||||
"TOTAL",
|
||||
all_latencies,
|
||||
total_errors,
|
||||
total_skips,
|
||||
total_bytes,
|
||||
wall_seconds,
|
||||
)
|
||||
)
|
||||
|
||||
headers = [
|
||||
"Action",
|
||||
"OK",
|
||||
"Err",
|
||||
"Skip",
|
||||
"Min ms",
|
||||
"Max ms",
|
||||
"Avg ms",
|
||||
"Med ms",
|
||||
"P95 ms",
|
||||
"P99 ms",
|
||||
"Std ms",
|
||||
"Ops/s",
|
||||
"MB",
|
||||
]
|
||||
fmt = "{:<12} {:>7} {:>5} {:>5} {:>9} {:>9} {:>9} {:>9} {:>9} {:>9} {:>9} {:>8} {:>9}"
|
||||
line = "-" * 122
|
||||
|
||||
print()
|
||||
print(line)
|
||||
print(f"Stress test report (wall time: {wall_seconds:.2f}s)")
|
||||
print(line)
|
||||
print(fmt.format(*headers))
|
||||
print(line)
|
||||
for r in rows:
|
||||
if r["action"] == "TOTAL":
|
||||
print(line)
|
||||
print(
|
||||
fmt.format(
|
||||
r["action"],
|
||||
r["count"],
|
||||
r["errors"],
|
||||
r["skips"],
|
||||
f"{r['min']:.2f}",
|
||||
f"{r['max']:.2f}",
|
||||
f"{r['avg']:.2f}",
|
||||
f"{r['median']:.2f}",
|
||||
f"{r['p95']:.2f}",
|
||||
f"{r['p99']:.2f}",
|
||||
f"{r['stddev']:.2f}",
|
||||
f"{r['ops']:.1f}",
|
||||
f"{r['mb']:.1f}",
|
||||
)
|
||||
)
|
||||
print(line)
|
||||
|
||||
|
||||
class Context:
|
||||
def __init__(self, args, stats):
|
||||
self.smtp_server = args.smtp_server
|
||||
self.smtp_port = args.smtp_port
|
||||
self.imap_server = args.imap_server
|
||||
self.imap_port = args.imap_port
|
||||
self.runs = args.runs
|
||||
self.max_content_size = args.max_content_size
|
||||
self.verbose = not args.quiet
|
||||
self.stats = stats
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Concurrent SMTP/IMAP stress test for Stalwart."
|
||||
)
|
||||
parser.add_argument("--smtp-server", default=DEFAULT_SMTP_SERVER)
|
||||
parser.add_argument("--smtp-port", type=int, default=DEFAULT_SMTP_PORT)
|
||||
parser.add_argument("--imap-server", default=DEFAULT_IMAP_SERVER)
|
||||
parser.add_argument("--imap-port", type=int, default=DEFAULT_IMAP_PORT)
|
||||
parser.add_argument("--threads", type=int, default=DEFAULT_THREADS)
|
||||
parser.add_argument(
|
||||
"--runs",
|
||||
type=int,
|
||||
default=DEFAULT_RUNS,
|
||||
help="Actions per thread. Use 0 for an infinite loop (stop with Ctrl-C).",
|
||||
)
|
||||
parser.add_argument("--credentials", default="users.txt")
|
||||
parser.add_argument(
|
||||
"--max-content-size",
|
||||
type=int,
|
||||
default=1048576,
|
||||
help="Maximum random message body size in bytes.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dict",
|
||||
default=DEFAULT_DICTIONARY,
|
||||
help="Word list used to generate message text, one word per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quiet",
|
||||
action="store_true",
|
||||
help="Suppress per-operation logging; print only the final report.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
global WORDS
|
||||
args = parse_args()
|
||||
if args.runs == 0:
|
||||
args.runs = None
|
||||
WORDS = load_words(args.dict)
|
||||
credentials = read_credentials(args.credentials)
|
||||
stats = Stats()
|
||||
ctx = Context(args, stats)
|
||||
|
||||
threads = [
|
||||
threading.Thread(target=thread_function, args=(ctx, credentials), daemon=True)
|
||||
for _ in range(args.threads)
|
||||
]
|
||||
|
||||
start = time.monotonic()
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
|
||||
try:
|
||||
while any(t.is_alive() for t in threads):
|
||||
for t in threads:
|
||||
t.join(timeout=0.2)
|
||||
except KeyboardInterrupt:
|
||||
print("\nStopping...")
|
||||
STOP_EVENT.set()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
wall_seconds = time.monotonic() - start
|
||||
print_report(stats, wall_seconds)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,244 @@
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import random
|
||||
import ssl
|
||||
import string
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
CORE = "urn:ietf:params:jmap:core"
|
||||
STALWART = "urn:stalwart:jmap"
|
||||
USING = [CORE, STALWART]
|
||||
|
||||
DEFAULT_BASE_URL = "https://127.0.0.1"
|
||||
DEFAULT_NUM_USERS = 1000
|
||||
DEFAULT_OUTPUT = "users.txt"
|
||||
DEFAULT_PASSWORD_LENGTH = 16
|
||||
DEFAULT_PREFIX = "test"
|
||||
|
||||
CREATE_RETRIES = 5
|
||||
|
||||
|
||||
class AccountError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def generate_password(length):
|
||||
return "".join(random.choices(string.ascii_letters + string.digits, k=length))
|
||||
|
||||
|
||||
def primary_account(session):
|
||||
accounts = session.get("primaryAccounts") or {}
|
||||
if "urn:ietf:params:jmap:mail" in accounts:
|
||||
return accounts["urn:ietf:params:jmap:mail"]
|
||||
all_accounts = session.get("accounts") or {}
|
||||
return next(iter(all_accounts), None)
|
||||
|
||||
|
||||
def build_opener(verify):
|
||||
context = ssl.create_default_context()
|
||||
if not verify:
|
||||
context.check_hostname = False
|
||||
context.verify_mode = ssl.CERT_NONE
|
||||
handler = urllib.request.HTTPSHandler(context=context)
|
||||
return urllib.request.build_opener(handler)
|
||||
|
||||
|
||||
class JmapClient:
|
||||
def __init__(self, base_url, auth_header, verify):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.auth_header = auth_header
|
||||
self.opener = build_opener(verify)
|
||||
self.api_url = None
|
||||
self.account_id = None
|
||||
self._discover()
|
||||
|
||||
def _http(self, url, method, body=None):
|
||||
headers = {"Authorization": self.auth_header}
|
||||
data = None
|
||||
if body is not None:
|
||||
data = json.dumps(body).encode("utf-8")
|
||||
headers["Content-Type"] = "application/json"
|
||||
request = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||||
try:
|
||||
with self.opener.open(request, timeout=60) as response:
|
||||
return response.status, response.read().decode("utf-8", "replace")
|
||||
except urllib.error.HTTPError as e:
|
||||
return e.code, e.read().decode("utf-8", "replace")
|
||||
except urllib.error.URLError as e:
|
||||
raise SystemExit(f"Request to {url} failed: {e.reason}")
|
||||
|
||||
def _discover(self):
|
||||
if "/.well-known/jmap" in self.base_url:
|
||||
candidates = [self.base_url]
|
||||
else:
|
||||
candidates = [self.base_url, f"{self.base_url}/.well-known/jmap"]
|
||||
last = ""
|
||||
for url in candidates:
|
||||
status, text = self._http(url, "GET")
|
||||
if status == 200:
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except ValueError:
|
||||
last = f"{url}: invalid JSON"
|
||||
continue
|
||||
if data.get("apiUrl") and data.get("accounts"):
|
||||
self.api_url = data["apiUrl"]
|
||||
self.account_id = primary_account(data)
|
||||
if not self.account_id:
|
||||
raise SystemExit(f"JMAP session at {url} has no accounts.")
|
||||
return
|
||||
last = f"{url}: status {status}"
|
||||
raise SystemExit(f"Could not discover JMAP session ({last}).")
|
||||
|
||||
def request(self, method_calls):
|
||||
body = {"using": USING, "methodCalls": method_calls}
|
||||
status, text = self._http(self.api_url, "POST", body)
|
||||
if status != 200:
|
||||
raise SystemExit(f"JMAP request failed: status {status}: {text}")
|
||||
return json.loads(text)
|
||||
|
||||
def call(self, method, args):
|
||||
args = dict(args)
|
||||
args["accountId"] = self.account_id
|
||||
parsed = self.request([[method, args, "c0"]])
|
||||
responses = parsed.get("methodResponses") or []
|
||||
if not responses:
|
||||
raise SystemExit(f"{method}: empty methodResponses")
|
||||
name, payload = responses[0][0], responses[0][1]
|
||||
if name == "error":
|
||||
raise SystemExit(f"{method} error: {payload}")
|
||||
return payload
|
||||
|
||||
def domain_id(self, name):
|
||||
calls = [
|
||||
[
|
||||
"x:Domain/query",
|
||||
{"accountId": self.account_id, "filter": {"name": name}},
|
||||
"q",
|
||||
],
|
||||
[
|
||||
"x:Domain/get",
|
||||
{
|
||||
"accountId": self.account_id,
|
||||
"#ids": {
|
||||
"resultOf": "q",
|
||||
"name": "x:Domain/query",
|
||||
"path": "/ids",
|
||||
},
|
||||
"properties": ["id", "name"],
|
||||
},
|
||||
"g",
|
||||
],
|
||||
]
|
||||
parsed = self.request(calls)
|
||||
responses = parsed.get("methodResponses") or []
|
||||
if len(responses) < 2:
|
||||
raise SystemExit(f"Domain lookup failed: {parsed}")
|
||||
get = responses[1][1]
|
||||
for entry in get.get("list") or []:
|
||||
if entry.get("name") == name:
|
||||
return entry.get("id")
|
||||
return None
|
||||
|
||||
def create_account(self, localpart, domain_id, password):
|
||||
create = {
|
||||
"a": {
|
||||
"@type": "User",
|
||||
"name": localpart,
|
||||
"domainId": domain_id,
|
||||
"credentials": {"0": {"@type": "Password", "secret": password}},
|
||||
"encryptionAtRest": {"@type": "Disabled"},
|
||||
"permissions": {"@type": "Inherit"},
|
||||
"roles": {"@type": "User"},
|
||||
"locale": "en_US",
|
||||
}
|
||||
}
|
||||
response = self.call("x:Account/set", {"create": create})
|
||||
not_created = response.get("notCreated") or {}
|
||||
if not_created:
|
||||
raise AccountError(f"{not_created.get('a', not_created)}")
|
||||
created = (response.get("created") or {}).get("a") or {}
|
||||
return created.get("id")
|
||||
|
||||
def invalidate_caches(self):
|
||||
self.call(
|
||||
"x:Action/set",
|
||||
{"create": {"c": {"@type": "InvalidateCaches"}}},
|
||||
)
|
||||
|
||||
|
||||
def build_auth_header(args):
|
||||
if args.token:
|
||||
return f"Bearer {args.token}"
|
||||
raw = f"{args.user}:{args.password}".encode("utf-8")
|
||||
return "Basic " + base64.b64encode(raw).decode("ascii")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Provision test accounts for the stress test using the JMAP API."
|
||||
)
|
||||
parser.add_argument("--base-url", default=DEFAULT_BASE_URL)
|
||||
parser.add_argument(
|
||||
"--domain",
|
||||
required=True,
|
||||
help="Existing domain name; resolved to a domain id via JMAP.",
|
||||
)
|
||||
parser.add_argument("--token", help="OAuth bearer token for the JMAP API.")
|
||||
parser.add_argument("--user", help="Basic auth username for the JMAP API.")
|
||||
parser.add_argument("--password", help="Basic auth password for the JMAP API.")
|
||||
parser.add_argument("--num-users", type=int, default=DEFAULT_NUM_USERS)
|
||||
parser.add_argument("--prefix", default=DEFAULT_PREFIX)
|
||||
parser.add_argument("--password-length", type=int, default=DEFAULT_PASSWORD_LENGTH)
|
||||
parser.add_argument("--output", default=DEFAULT_OUTPUT)
|
||||
parser.add_argument(
|
||||
"--verify",
|
||||
action="store_true",
|
||||
help="Verify TLS certificates (disabled by default for self-signed servers).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if not args.token and not (args.user and args.password):
|
||||
parser.error("provide either --token or both --user and --password")
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
client = JmapClient(args.base_url, build_auth_header(args), args.verify)
|
||||
domain_id = client.domain_id(args.domain)
|
||||
if not domain_id:
|
||||
raise SystemExit(f"Domain '{args.domain}' not found via JMAP.")
|
||||
|
||||
created = 0
|
||||
failed = 0
|
||||
with open(args.output, "w") as file:
|
||||
for i in range(1, args.num_users + 1):
|
||||
localpart = f"{args.prefix}{i}"
|
||||
email = f"{localpart}@{args.domain}"
|
||||
password = None
|
||||
last_error = None
|
||||
for _ in range(CREATE_RETRIES):
|
||||
password = generate_password(args.password_length)
|
||||
try:
|
||||
client.create_account(localpart, domain_id, password)
|
||||
last_error = None
|
||||
break
|
||||
except AccountError as e:
|
||||
last_error = e
|
||||
if last_error is not None:
|
||||
failed += 1
|
||||
print(f"FAIL {email}: {last_error}")
|
||||
continue
|
||||
file.write(f"{email}:{password}\n")
|
||||
file.flush()
|
||||
created += 1
|
||||
print(f"OK {email}")
|
||||
|
||||
client.invalidate_caches()
|
||||
print(f"\nCreated {created} accounts ({failed} failed). Written to {args.output}.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user