553 lines
26 KiB
Python
553 lines
26 KiB
Python
import socketserver
|
|
import threading
|
|
|
|
RESPONSE_SEARCH_COMPLETED = "OK SEARCH completed"
|
|
RESPONSE_SELECT_FIRST = "NO Select first"
|
|
|
|
|
|
class MockIMAPHandler(socketserver.StreamRequestHandler):
|
|
"""
|
|
A minimal IMAP4rev1 mock server handler for testing purposes.
|
|
Supports basic commands required for the migration scripts.
|
|
"""
|
|
|
|
def handle(self):
|
|
self.wfile.write(b"* OK [CAPABILITY IMAP4rev1] Mock IMAP Server Ready\r\n")
|
|
self.selected_folder = None
|
|
|
|
# Check if server has folders, if not, initialize
|
|
if hasattr(self.server, "folders"):
|
|
self.current_folders = self.server.folders
|
|
else:
|
|
# Just in case it's used standalone without the threaded wrapper
|
|
self.current_folders = {"INBOX": []}
|
|
|
|
self.retry_state = {} # key -> count
|
|
|
|
while True:
|
|
try:
|
|
line = self.rfile.readline()
|
|
if not line:
|
|
break
|
|
line = line.decode("utf-8").strip()
|
|
if not line:
|
|
continue
|
|
|
|
parts = line.split(" ", 2)
|
|
tag = parts[0]
|
|
cmd = parts[1].upper()
|
|
args = parts[2] if len(parts) > 2 else ""
|
|
|
|
# Check for RETRY_N injection in args (for testing imap_retry)
|
|
# Pattern: RETRY_N where N is integer
|
|
import re
|
|
|
|
retry_match = re.search(r"RETRY_(\d+)", args)
|
|
if retry_match:
|
|
count = int(retry_match.group(1))
|
|
# Create a unique key for this command invocation context
|
|
# Just using command + args is roughly sufficient
|
|
key = f"{cmd}_{args}"
|
|
current_fails = self.retry_state.get(key, 0)
|
|
|
|
if current_fails < count:
|
|
# Logic Fix: Only increment if we are going to fail
|
|
self.retry_state[key] = current_fails + 1
|
|
self.send_response(tag, "NO [UNAVAILABLE] Server Busy (Simulated)")
|
|
continue
|
|
# Else fall through to normal command processing
|
|
|
|
if cmd == "LOGIN":
|
|
self.send_response(tag, "OK LOGIN completed")
|
|
|
|
elif cmd == "APPEND":
|
|
# Check for literal size
|
|
import re
|
|
|
|
size_match = re.search(r"\{(\d+)\}$", args)
|
|
data = b""
|
|
if size_match:
|
|
size = int(size_match.group(1))
|
|
# Send continuation
|
|
self.wfile.write(b"+\r\n")
|
|
self.wfile.flush()
|
|
# Read data
|
|
data = self.rfile.read(size)
|
|
|
|
# Extract mailbox name to store the message
|
|
# Simplistic parsing: check for quoted or unquoted first argument
|
|
mailbox = "INBOX" # Default fallback
|
|
clean_args = args.lstrip()
|
|
if clean_args.startswith('"'):
|
|
end_quote = clean_args.find('"', 1)
|
|
if end_quote != -1:
|
|
mailbox = clean_args[1:end_quote]
|
|
else:
|
|
mailbox = clean_args.split(" ")[0]
|
|
|
|
if mailbox not in self.current_folders:
|
|
self.current_folders[mailbox] = []
|
|
|
|
# Parse flags if present: look for (...)
|
|
msg_flags = set()
|
|
flag_match = re.search(r"\(([^)]+)\)", args)
|
|
if flag_match:
|
|
# e.g. (\Seen \Deleted)
|
|
flag_content = flag_match.group(1)
|
|
msg_flags = set(flag_content.split())
|
|
|
|
# Store message
|
|
new_uid = len(self.current_folders[mailbox]) + 1
|
|
msg_obj = {"uid": new_uid, "flags": msg_flags, "content": data}
|
|
self.current_folders[mailbox].append(msg_obj)
|
|
|
|
# Always succeed for test if we passed retry logic
|
|
self.wfile.write(tag.encode() + b" OK [APPENDUID 1 100] APPEND completed\r\n")
|
|
self.wfile.flush()
|
|
|
|
elif cmd == "LOGOUT":
|
|
self.send_response(tag, "OK LOGOUT completed")
|
|
break
|
|
|
|
elif cmd == "CAPABILITY":
|
|
self.wfile.write(b"* CAPABILITY IMAP4rev1 AUTH=PLAIN\r\n")
|
|
self.send_response(tag, "OK CAPABILITY completed")
|
|
|
|
elif cmd == "AUTHENTICATE":
|
|
# Minimal XOAUTH2 support for tests.
|
|
self.wfile.write(b"+ \r\n")
|
|
_ = self.rfile.readline()
|
|
# Also consume potential retry line? No.
|
|
self.send_response(tag, "OK AUTHENTICATE completed")
|
|
|
|
elif cmd == "NOOP":
|
|
self.send_response(tag, "OK NOOP completed")
|
|
|
|
elif cmd == "LIST":
|
|
for folder in self.current_folders:
|
|
self.wfile.write(f'* LIST (\\HasNoChildren) "/" "{folder}"\r\n'.encode())
|
|
self.send_response(tag, "OK LIST completed")
|
|
|
|
elif cmd == "SELECT":
|
|
folder = args.strip().strip('"')
|
|
|
|
# For retry testing: if folder starts with NO, fail immediately
|
|
if folder.startswith("NO"):
|
|
# e.g. "NO [UNAVAILABLE]"
|
|
self.send_response(tag, folder)
|
|
continue
|
|
|
|
# Allow RETRY_N folders to succeed if they passed the retry intercept logic
|
|
is_retry_folder = "RETRY_" in folder
|
|
|
|
if folder in self.current_folders or is_retry_folder:
|
|
self.selected_folder = folder
|
|
count = len(self.current_folders.get(folder, [])) # Default empty for retry folders
|
|
self.wfile.write(f"* {count} EXISTS\r\n".encode())
|
|
self.wfile.write(f"* {count} RECENT\r\n".encode())
|
|
self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
|
|
self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n")
|
|
self.send_response(tag, "OK [READ-WRITE] SELECT completed")
|
|
elif folder == "EMPTY":
|
|
# Special case for "empty_data" test in retry
|
|
self.wfile.write(b"* 0 EXISTS\r\n")
|
|
self.send_response(tag, "OK SELECT completed")
|
|
else:
|
|
self.send_response(tag, "NO [NONEXISTENT] Folder not found")
|
|
|
|
elif cmd == "EXAMINE":
|
|
# EXAMINE is like SELECT but read-only
|
|
folder = args.strip().strip('"')
|
|
|
|
is_retry_folder = "RETRY_" in folder
|
|
|
|
if folder in self.current_folders or is_retry_folder:
|
|
self.selected_folder = folder
|
|
count = len(self.current_folders.get(folder, []))
|
|
self.wfile.write(f"* {count} EXISTS\r\n".encode())
|
|
self.wfile.write(f"* {count} RECENT\r\n".encode())
|
|
self.wfile.write(b"* FLAGS (\\Seen \\Answered \\Flagged \\Deleted \\Draft)\r\n")
|
|
self.wfile.write(b"* OK [UIDVALIDITY 1] UIDs valid\r\n")
|
|
self.send_response(tag, "OK [READ-ONLY] EXAMINE completed")
|
|
else:
|
|
self.send_response(tag, "NO [NONEXISTENT] Folder not found")
|
|
|
|
elif cmd == "CREATE":
|
|
folder = args.strip().strip('"')
|
|
if folder not in self.current_folders:
|
|
self.current_folders[folder] = []
|
|
self.send_response(tag, "OK CREATE completed")
|
|
|
|
elif cmd == "SEARCH":
|
|
if not self.selected_folder:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
continue
|
|
|
|
msgs = self.current_folders[self.selected_folder]
|
|
sub_args = args
|
|
|
|
header_msg_id = None
|
|
try:
|
|
m = re.search(r'HEADER\s+Message-ID\s+"([^"]+)"', sub_args, re.IGNORECASE)
|
|
if m:
|
|
header_msg_id = m.group(1)
|
|
except Exception:
|
|
header_msg_id = None
|
|
|
|
seq_nums = []
|
|
for idx, m in enumerate(msgs, start=1):
|
|
if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]:
|
|
continue
|
|
if header_msg_id:
|
|
msg_text = m["content"].decode("utf-8", errors="ignore")
|
|
if header_msg_id not in msg_text:
|
|
continue
|
|
seq_nums.append(str(idx))
|
|
|
|
if "ALL" in sub_args.upper() and not header_msg_id:
|
|
seq_nums = [str(idx) for idx in range(1, len(msgs) + 1)]
|
|
|
|
seq_str = " ".join(seq_nums)
|
|
if seq_str:
|
|
self.wfile.write(f"* SEARCH {seq_str}\r\n".encode())
|
|
else:
|
|
self.wfile.write(b"* SEARCH\r\n")
|
|
self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
|
|
|
|
elif cmd == "EXPUNGE":
|
|
if self.selected_folder:
|
|
msgs = self.current_folders[self.selected_folder]
|
|
# In-place remove marked deleted
|
|
# We use a dictionary-like structure implicitly via dicts in list now
|
|
# Check "flags" set in msg object
|
|
new_msgs = [m for m in msgs if "\\Deleted" not in m["flags"]]
|
|
self.current_folders[self.selected_folder] = new_msgs
|
|
# According to IMAP, we should send Expunge indices but we'll skip for mock
|
|
self.send_response(tag, "OK EXPUNGE completed")
|
|
else:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
|
|
elif cmd == "UID":
|
|
sub_parts = args.split(" ", 1)
|
|
sub_cmd = sub_parts[0].upper()
|
|
sub_rest = sub_parts[1] if len(sub_parts) > 1 else ""
|
|
|
|
if sub_cmd == "SEARCH":
|
|
sub_args = sub_rest
|
|
if not self.selected_folder:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
continue
|
|
msgs = self.current_folders[self.selected_folder]
|
|
|
|
# Support UID range searches like: UID SEARCH UID 101:*
|
|
uid_min = None
|
|
try:
|
|
m = re.search(r"UID\s+(\d+)", sub_args, re.IGNORECASE)
|
|
if m:
|
|
uid_min = int(m.group(1))
|
|
except Exception:
|
|
uid_min = None
|
|
|
|
# Handle UNDELETED
|
|
# If args has UNDELETED, filter flags
|
|
valid_uids = []
|
|
for m in msgs:
|
|
if uid_min is not None and m["uid"] < uid_min:
|
|
continue
|
|
if "UNDELETED" in sub_args and "\\Deleted" in m["flags"]:
|
|
continue
|
|
valid_uids.append(str(m["uid"]))
|
|
|
|
uids_str = " ".join(valid_uids)
|
|
if uids_str:
|
|
self.wfile.write(f"* SEARCH {uids_str}\r\n".encode())
|
|
else:
|
|
self.wfile.write(b"* SEARCH\r\n")
|
|
self.send_response(tag, RESPONSE_SEARCH_COMPLETED)
|
|
|
|
elif sub_cmd == "STORE":
|
|
# UID STORE <uid> +FLAGS (\Deleted)
|
|
store_parts = sub_rest.split(" ", 2)
|
|
|
|
uid = int(store_parts[0])
|
|
action = store_parts[1].upper() # +FLAGS or -FLAGS
|
|
|
|
# Flags can be "(\Deleted)" or "\Deleted"
|
|
flags_str = store_parts[2].strip("()")
|
|
flags_list = {f.strip() for f in flags_str.split()}
|
|
|
|
if self.selected_folder:
|
|
msgs = self.current_folders[self.selected_folder]
|
|
found = False
|
|
for m in msgs:
|
|
if m["uid"] == uid:
|
|
if action == "+FLAGS":
|
|
m["flags"].update(flags_list)
|
|
elif action == "-FLAGS":
|
|
m["flags"].difference_update(flags_list)
|
|
found = True
|
|
|
|
# Send update
|
|
flag_output = " ".join(m["flags"])
|
|
self.wfile.write(
|
|
f"* {msgs.index(m) + 1} FETCH (FLAGS ({flag_output}))\r\n".encode()
|
|
)
|
|
break
|
|
if found:
|
|
self.send_response(tag, "OK STORE completed")
|
|
else:
|
|
self.send_response(tag, "NO UID not found")
|
|
else:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
|
|
elif sub_cmd == "FETCH":
|
|
# sub_rest is e.g. "1 (RFC822.SIZE BODY.PEEK[...])"
|
|
parts = sub_rest.split(" ", 1)
|
|
uid_set = parts[0]
|
|
opts = parts[1].upper() if len(parts) > 1 else ""
|
|
|
|
if not self.selected_folder:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
continue
|
|
|
|
msgs = self.current_folders[self.selected_folder]
|
|
# Find msg by UID - handle single UID, comma-separated, or ranges
|
|
target_msgs = []
|
|
if ":" in uid_set:
|
|
# Range logic omitted for brevity, taking all for '*' or simplified assumption
|
|
# But since we use UIDs, we should filter.
|
|
target_msgs = msgs # Return all for range
|
|
elif "," in uid_set:
|
|
# Handle comma-separated UIDs like "1,2,3"
|
|
try:
|
|
uid_list = [int(u) for u in uid_set.split(",")]
|
|
target_msgs = [m for m in msgs if m["uid"] in uid_list]
|
|
except Exception:
|
|
pass
|
|
else:
|
|
try:
|
|
t_uid = int(uid_set)
|
|
target_msgs = [m for m in msgs if m["uid"] == t_uid]
|
|
except Exception:
|
|
pass
|
|
|
|
for m in target_msgs:
|
|
# Mock always returns info if iterating all, or specific
|
|
# Construction
|
|
msg_content = m["content"]
|
|
msg_len = len(msg_content)
|
|
flags_str = " ".join(m["flags"])
|
|
|
|
# Check if requesting header fields only
|
|
if "HEADER.FIELDS" in opts:
|
|
rfc822_size_part = f" RFC822.SIZE {msg_len}" if "RFC822.SIZE" in opts else ""
|
|
# Extract just the headers from content
|
|
content_str = msg_content.decode("utf-8", errors="ignore")
|
|
header_part = content_str.split("\r\n\r\n")[0] + "\r\n"
|
|
header_bytes = header_part.encode("utf-8")
|
|
header_len = len(header_bytes)
|
|
|
|
resp = (
|
|
f"* {msgs.index(m) + 1} FETCH (UID {m['uid']}{rfc822_size_part} FLAGS ({flags_str}) "
|
|
f"BODY[HEADER.FIELDS (MESSAGE-ID SUBJECT)] {{{header_len}}}\r\n"
|
|
)
|
|
self.wfile.write(resp.encode("utf-8"))
|
|
self.wfile.write(header_bytes)
|
|
self.wfile.write(b")\r\n")
|
|
elif "RFC822" in opts or "BODY" in opts:
|
|
resp = f"* {msgs.index(m) + 1} FETCH (UID {m['uid']} RFC822.SIZE {msg_len} FLAGS ({flags_str})"
|
|
resp += f" BODY[] {{{msg_len}}}\r\n"
|
|
self.wfile.write(resp.encode("utf-8"))
|
|
self.wfile.write(msg_content)
|
|
self.wfile.write(b")\r\n")
|
|
else:
|
|
resp = f"* {msgs.index(m) + 1} FETCH (UID {m['uid']} RFC822.SIZE {msg_len} FLAGS ({flags_str})"
|
|
resp += ")\r\n"
|
|
self.wfile.write(resp.encode("utf-8"))
|
|
|
|
self.wfile.flush()
|
|
|
|
self.send_response(tag, "OK FETCH completed")
|
|
|
|
elif sub_cmd == "COPY":
|
|
# UID COPY <uid> <target>
|
|
c_parts = sub_rest.split(" ", 1)
|
|
c_uid = int(c_parts[0])
|
|
c_dest = c_parts[1].strip().strip('"')
|
|
|
|
if c_dest in self.current_folders and self.selected_folder:
|
|
msgs = self.current_folders[self.selected_folder]
|
|
found_msg = next((m for m in msgs if m["uid"] == c_uid), None)
|
|
if found_msg:
|
|
# COPY: create new msg object (new UID usually)
|
|
# We need a new UID generator.
|
|
# For simplicity, use max_uid + 1
|
|
dest_msgs = self.current_folders[c_dest]
|
|
max_uid = max([m["uid"] for m in dest_msgs], default=0)
|
|
new_msg = found_msg.copy()
|
|
new_msg["uid"] = max_uid + 1
|
|
new_msg["flags"] = found_msg["flags"].copy()
|
|
dest_msgs.append(new_msg)
|
|
self.send_response(tag, "OK COPY completed")
|
|
else:
|
|
self.send_response(tag, "NO UID not found")
|
|
else:
|
|
self.send_response(tag, "NO Dest not found")
|
|
|
|
elif cmd == "APPEND":
|
|
try:
|
|
match = re.search(r"\{(\d+)\}$", args)
|
|
if match:
|
|
size = int(match.group(1))
|
|
self.wfile.write(b"+ Ready\r\n")
|
|
data = self.rfile.read(size)
|
|
|
|
# Args format (typical): <folder> (<flags>) "<internaldate>" {<size>}
|
|
# We only care about folder + flags for tests.
|
|
folder_arg = args.split(" ")[0].strip().strip('"')
|
|
flags_match = re.search(r"\(([^)]*)\)", args)
|
|
flags_set = set()
|
|
if flags_match:
|
|
flags_str = flags_match.group(1).strip()
|
|
if flags_str:
|
|
flags_set = {f.strip() for f in flags_str.split() if f.strip()}
|
|
|
|
if folder_arg in self.current_folders:
|
|
dest_msgs = self.current_folders[folder_arg]
|
|
max_uid = max([m["uid"] for m in dest_msgs], default=0)
|
|
|
|
# IMPORTANT: Reset pointer or copy
|
|
new_msg = {"uid": max_uid + 1, "flags": flags_set, "content": data}
|
|
dest_msgs.append(new_msg)
|
|
print(f"MOCK APPEND SUCCESS: {folder_arg} now has {len(dest_msgs)}")
|
|
self.send_response(tag, "OK APPEND completed")
|
|
else:
|
|
self.send_response(tag, "NO Folder not found")
|
|
else:
|
|
self.send_response(tag, "BAD APPEND")
|
|
except Exception as e:
|
|
print(f"MOCK APPEND ERROR: {e}")
|
|
self.send_response(tag, "BAD APPEND")
|
|
|
|
elif cmd == "STORE":
|
|
# STORE <msg_set> +FLAGS (\Seen) (non-UID; uses message sequence numbers)
|
|
if not self.selected_folder:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
continue
|
|
|
|
try:
|
|
store_parts = args.split(" ", 2)
|
|
msg_set = store_parts[0]
|
|
action = store_parts[1].upper() # +FLAGS or -FLAGS
|
|
flags_str = store_parts[2].strip().strip("()")
|
|
flags_list = {f.strip() for f in flags_str.split() if f.strip()}
|
|
|
|
msgs = self.current_folders[self.selected_folder]
|
|
|
|
# Only implement single message number for now (sufficient for tests)
|
|
try:
|
|
msg_num = int(msg_set)
|
|
except Exception:
|
|
self.send_response(tag, "BAD STORE")
|
|
continue
|
|
|
|
if msg_num < 1 or msg_num > len(msgs):
|
|
self.send_response(tag, "NO Message not found")
|
|
continue
|
|
|
|
m = msgs[msg_num - 1]
|
|
if action == "+FLAGS":
|
|
m["flags"].update(flags_list)
|
|
elif action == "-FLAGS":
|
|
m["flags"].difference_update(flags_list)
|
|
|
|
flag_output = " ".join(m["flags"])
|
|
self.wfile.write(f"* {msg_num} FETCH (FLAGS ({flag_output}))\r\n".encode())
|
|
self.send_response(tag, "OK STORE completed")
|
|
except Exception:
|
|
self.send_response(tag, "BAD STORE")
|
|
|
|
elif cmd == "FETCH":
|
|
# Non-UID FETCH - args is e.g. "1 (RFC822.SIZE)"
|
|
if not self.selected_folder:
|
|
self.send_response(tag, RESPONSE_SELECT_FIRST)
|
|
continue
|
|
|
|
parts = args.split(" ", 1)
|
|
msg_num = int(parts[0])
|
|
opts = parts[1].upper() if len(parts) > 1 else ""
|
|
|
|
msgs = self.current_folders[self.selected_folder]
|
|
if 1 <= msg_num <= len(msgs):
|
|
m = msgs[msg_num - 1] # Convert to 0-indexed
|
|
msg_content = m["content"]
|
|
msg_len = len(msg_content)
|
|
flags_str = " ".join(m["flags"])
|
|
|
|
resp = f"* {msg_num} FETCH (UID {m['uid']} RFC822.SIZE {msg_len} FLAGS ({flags_str})"
|
|
|
|
if "RFC822" in opts or "BODY" in opts:
|
|
resp += f" BODY[] {{{msg_len}}}\r\n"
|
|
self.wfile.write(resp.encode("utf-8"))
|
|
self.wfile.write(msg_content)
|
|
self.wfile.write(b")\r\n")
|
|
else:
|
|
resp += ")\r\n"
|
|
self.wfile.write(resp.encode("utf-8"))
|
|
self.wfile.flush()
|
|
|
|
self.send_response(tag, "OK FETCH completed")
|
|
|
|
elif cmd == "NOOP":
|
|
self.send_response(tag, "OK NOOP")
|
|
|
|
else:
|
|
self.send_response(tag, "BAD Command not recognized")
|
|
|
|
except Exception as e:
|
|
print(f"MockServer EXCEPTION: {e}")
|
|
import traceback
|
|
|
|
traceback.print_exc()
|
|
break
|
|
|
|
def send_response(self, tag, message):
|
|
self.wfile.write(f"{tag} {message}\r\n".encode())
|
|
|
|
|
|
class MockIMAPServer(socketserver.ThreadingMixIn, socketserver.TCPServer):
|
|
allow_reuse_address = True
|
|
daemon_threads = True
|
|
|
|
def __init__(self, server_address, request_handler_class, initial_folders=None):
|
|
super().__init__(server_address, request_handler_class)
|
|
self.folders = {}
|
|
if initial_folders:
|
|
for fname, contents in initial_folders.items():
|
|
self.folders[fname] = []
|
|
for i, c in enumerate(contents):
|
|
if isinstance(c, bytes):
|
|
self.folders[fname].append({"uid": i + 1, "flags": set(), "content": c})
|
|
else:
|
|
self.folders[fname].append(c)
|
|
else:
|
|
self.folders = {"INBOX": []}
|
|
|
|
|
|
def start_server_thread(port=0, initial_folders=None):
|
|
# Use port 0 to let OS select free port
|
|
if isinstance(port, dict):
|
|
# Handle legacy call where port was inadvertently passed as dict
|
|
initial_folders = port
|
|
port = 0
|
|
|
|
server = MockIMAPServer(("localhost", port), MockIMAPHandler, initial_folders)
|
|
|
|
# Get the actual port if 0 was used
|
|
actual_port = server.server_address[1]
|
|
|
|
t = threading.Thread(target=server.serve_forever)
|
|
t.daemon = True
|
|
t.start()
|
|
return server, actual_port
|
