Add progress bars, fix DCT decode, sparkly MOTD
Progress bar support (v4.1.2): - Web frontend: Real-time progress during encode with phase display - CLI: --progress flag with rich library for encode command - Backend: progress_file parameter for async progress reporting DCT decode bug fix: - Fixed InvalidMagicBytesError not being caught in early-exit check - RS-protected format (v4.1.0+) has length prefix first, not magic bytes - Exception handler now catches both ValueError and InvalidMagicBytesError MOTD update: - Added sparkly header to setup.sh MOTD (matches other rpi scripts) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
@@ -24,11 +24,31 @@ Usage:
|
||||
stegasoo channel [SUBCOMMAND]
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import click
|
||||
|
||||
# Rich progress bar (optional)
|
||||
try:
|
||||
from rich.progress import (
|
||||
BarColumn,
|
||||
Progress,
|
||||
SpinnerColumn,
|
||||
TaskProgressColumn,
|
||||
TextColumn,
|
||||
TimeElapsedColumn,
|
||||
)
|
||||
|
||||
HAS_RICH = True
|
||||
except ImportError:
|
||||
HAS_RICH = False
|
||||
|
||||
# Add parent to path for development
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
|
||||
|
||||
@@ -598,6 +618,73 @@ def channel_clear(project, clear_all, force):
|
||||
click.echo(" Mode is now: PUBLIC")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# PROGRESS BAR UTILITIES (v4.1.2)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _generate_progress_job_id() -> str:
|
||||
"""Generate a unique job ID for progress tracking."""
|
||||
return str(uuid.uuid4())[:8]
|
||||
|
||||
|
||||
def _get_progress_file_path(job_id: str) -> str:
|
||||
"""Get the progress file path for a job ID."""
|
||||
return str(Path(tempfile.gettempdir()) / f"stegasoo_progress_{job_id}.json")
|
||||
|
||||
|
||||
def _read_progress(job_id: str) -> dict | None:
|
||||
"""Read progress from file for a job ID."""
|
||||
progress_file = _get_progress_file_path(job_id)
|
||||
try:
|
||||
with open(progress_file) as f:
|
||||
return json.load(f)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def _cleanup_progress_file(job_id: str) -> None:
|
||||
"""Remove progress file for a completed job."""
|
||||
progress_file = _get_progress_file_path(job_id)
|
||||
try:
|
||||
Path(progress_file).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _run_encode_with_progress(encode_func, encode_kwargs: dict, progress_file: str) -> tuple:
|
||||
"""
|
||||
Run encode in a thread and return result.
|
||||
|
||||
Returns:
|
||||
(success, result_or_error)
|
||||
"""
|
||||
result_holder = {"result": None, "error": None}
|
||||
|
||||
def run():
|
||||
try:
|
||||
result_holder["result"] = encode_func(**encode_kwargs, progress_file=progress_file)
|
||||
except Exception as e:
|
||||
result_holder["error"] = e
|
||||
|
||||
thread = threading.Thread(target=run)
|
||||
thread.start()
|
||||
return thread, result_holder
|
||||
|
||||
|
||||
def _format_phase(phase: str) -> str:
|
||||
"""Format phase name for display."""
|
||||
phases = {
|
||||
"starting": "Starting",
|
||||
"initializing": "Initializing",
|
||||
"embedding": "Embedding",
|
||||
"saving": "Saving",
|
||||
"finalizing": "Finalizing",
|
||||
"complete": "Complete",
|
||||
}
|
||||
return phases.get(phase, phase.capitalize())
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ENCODE COMMAND
|
||||
# ============================================================================
|
||||
@@ -642,6 +729,7 @@ def channel_clear(project, clear_all, force):
|
||||
help="DCT color mode: grayscale (default) or color (preserves original colors)",
|
||||
)
|
||||
@click.option("--quiet", "-q", is_flag=True, help="Suppress output except errors")
|
||||
@click.option("--progress", is_flag=True, help="Show progress bar (requires rich)")
|
||||
def encode_cmd(
|
||||
ref,
|
||||
carrier,
|
||||
@@ -661,6 +749,7 @@ def encode_cmd(
|
||||
dct_output_format,
|
||||
dct_color_mode,
|
||||
quiet,
|
||||
progress,
|
||||
):
|
||||
"""
|
||||
Encode a secret message or file into an image.
|
||||
@@ -808,19 +897,63 @@ def encode_cmd(
|
||||
click.echo(channel_status)
|
||||
|
||||
# v4.0.0: Include channel_key parameter
|
||||
result = encode(
|
||||
message=payload,
|
||||
reference_photo=ref_photo,
|
||||
carrier_image=carrier_image,
|
||||
passphrase=passphrase,
|
||||
pin=pin or "",
|
||||
rsa_key_data=rsa_key_data,
|
||||
rsa_password=effective_key_password,
|
||||
embed_mode=embed_mode,
|
||||
dct_output_format=dct_output_format,
|
||||
dct_color_mode=dct_color_mode,
|
||||
channel_key=resolved_channel_key,
|
||||
)
|
||||
# v4.1.2: Progress bar support
|
||||
encode_kwargs = {
|
||||
"message": payload,
|
||||
"reference_photo": ref_photo,
|
||||
"carrier_image": carrier_image,
|
||||
"passphrase": passphrase,
|
||||
"pin": pin or "",
|
||||
"rsa_key_data": rsa_key_data,
|
||||
"rsa_password": effective_key_password,
|
||||
"embed_mode": embed_mode,
|
||||
"dct_output_format": dct_output_format,
|
||||
"dct_color_mode": dct_color_mode,
|
||||
"channel_key": resolved_channel_key,
|
||||
}
|
||||
|
||||
if progress and HAS_RICH:
|
||||
# Run with progress bar
|
||||
job_id = _generate_progress_job_id()
|
||||
progress_file = _get_progress_file_path(job_id)
|
||||
|
||||
thread, result_holder = _run_encode_with_progress(encode, encode_kwargs, progress_file)
|
||||
|
||||
with Progress(
|
||||
SpinnerColumn(),
|
||||
TextColumn("[progress.description]{task.description}"),
|
||||
BarColumn(),
|
||||
TaskProgressColumn(),
|
||||
TimeElapsedColumn(),
|
||||
transient=True,
|
||||
) as progress_bar:
|
||||
task = progress_bar.add_task("Encoding...", total=100)
|
||||
|
||||
while thread.is_alive():
|
||||
prog = _read_progress(job_id)
|
||||
if prog:
|
||||
percent = prog.get("percent", 0)
|
||||
phase = _format_phase(prog.get("phase", "processing"))
|
||||
progress_bar.update(task, completed=percent, description=f"{phase}...")
|
||||
time.sleep(0.1)
|
||||
|
||||
# Final update
|
||||
progress_bar.update(task, completed=100, description="Complete!")
|
||||
|
||||
_cleanup_progress_file(job_id)
|
||||
|
||||
if result_holder["error"]:
|
||||
raise result_holder["error"]
|
||||
result = result_holder["result"]
|
||||
|
||||
elif progress and not HAS_RICH:
|
||||
click.secho(
|
||||
"Warning: --progress requires 'rich' package. Install with: pip install rich",
|
||||
fg="yellow",
|
||||
)
|
||||
result = encode(**encode_kwargs)
|
||||
else:
|
||||
result = encode(**encode_kwargs)
|
||||
|
||||
# Determine output path
|
||||
if output:
|
||||
|
||||
@@ -26,7 +26,9 @@ import mimetypes
|
||||
import os
|
||||
import secrets
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from auth import (
|
||||
@@ -36,6 +38,7 @@ from auth import (
|
||||
can_create_user,
|
||||
can_save_channel_key,
|
||||
change_password,
|
||||
clear_recovery_key,
|
||||
create_admin_user,
|
||||
create_user,
|
||||
delete_channel_key,
|
||||
@@ -45,12 +48,11 @@ from auth import (
|
||||
get_channel_key_by_id,
|
||||
get_current_user,
|
||||
get_non_admin_count,
|
||||
get_recovery_key_hash,
|
||||
get_user_by_id,
|
||||
get_user_channel_keys,
|
||||
get_username,
|
||||
has_recovery_key,
|
||||
get_recovery_key_hash,
|
||||
clear_recovery_key,
|
||||
is_admin,
|
||||
is_authenticated,
|
||||
login_required,
|
||||
@@ -59,10 +61,10 @@ from auth import (
|
||||
reset_user_password,
|
||||
save_channel_key,
|
||||
set_recovery_key_hash,
|
||||
verify_and_reset_admin_password,
|
||||
update_channel_key_last_used,
|
||||
update_channel_key_name,
|
||||
user_exists,
|
||||
verify_and_reset_admin_password,
|
||||
verify_user_password,
|
||||
)
|
||||
from auth import (
|
||||
@@ -156,7 +158,13 @@ except ImportError:
|
||||
# ============================================================================
|
||||
# Runs encode/decode/compare in subprocesses to prevent jpegio/scipy crashes
|
||||
# from taking down the Flask server.
|
||||
from subprocess_stego import SubprocessStego
|
||||
from subprocess_stego import (
|
||||
SubprocessStego,
|
||||
cleanup_progress_file,
|
||||
generate_job_id,
|
||||
get_progress_file_path,
|
||||
read_progress,
|
||||
)
|
||||
|
||||
from stegasoo.qr_utils import (
|
||||
can_fit_in_qr,
|
||||
@@ -195,6 +203,42 @@ app.config["HTTPS_ENABLED"] = os.environ.get("STEGASOO_HTTPS_ENABLED", "false").
|
||||
# Initialize auth module
|
||||
init_auth(app)
|
||||
|
||||
# ============================================================================
|
||||
# ASYNC JOB MANAGEMENT (v4.1.2)
|
||||
# ============================================================================
|
||||
# Encode operations can run in background threads with progress reporting
|
||||
|
||||
# Thread pool for background encode/decode operations
|
||||
_executor = ThreadPoolExecutor(max_workers=2)
|
||||
|
||||
# Job storage: job_id -> {status, result, error, file_id, ...}
|
||||
_jobs = {}
|
||||
_jobs_lock = threading.Lock()
|
||||
|
||||
|
||||
def _store_job(job_id: str, data: dict) -> None:
|
||||
"""Thread-safe job storage."""
|
||||
with _jobs_lock:
|
||||
_jobs[job_id] = data
|
||||
|
||||
|
||||
def _get_job(job_id: str) -> dict | None:
|
||||
"""Thread-safe job retrieval."""
|
||||
with _jobs_lock:
|
||||
return _jobs.get(job_id)
|
||||
|
||||
|
||||
def _cleanup_old_jobs(max_age_seconds: int = 3600) -> None:
|
||||
"""Remove jobs older than max_age_seconds."""
|
||||
now = time.time()
|
||||
with _jobs_lock:
|
||||
to_remove = [
|
||||
jid for jid, data in _jobs.items() if now - data.get("created", 0) > max_age_seconds
|
||||
]
|
||||
for jid in to_remove:
|
||||
cleanup_progress_file(jid)
|
||||
del _jobs[jid]
|
||||
|
||||
|
||||
@app.before_request
|
||||
def require_setup():
|
||||
@@ -817,10 +861,119 @@ def api_check_fit():
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _run_encode_job(job_id: str, encode_params: dict) -> None:
|
||||
"""Background thread function for async encode."""
|
||||
progress_file = get_progress_file_path(job_id)
|
||||
|
||||
try:
|
||||
_store_job(job_id, {"status": "running", "created": time.time()})
|
||||
|
||||
# Run encode with progress file
|
||||
if encode_params.get("file_data"):
|
||||
encode_result = subprocess_stego.encode(
|
||||
carrier_data=encode_params["carrier_data"],
|
||||
reference_data=encode_params["ref_data"],
|
||||
file_data=encode_params["file_data"],
|
||||
file_name=encode_params["file_name"],
|
||||
file_mime=encode_params["file_mime"],
|
||||
passphrase=encode_params["passphrase"],
|
||||
pin=encode_params.get("pin"),
|
||||
rsa_key_data=encode_params.get("rsa_key_data"),
|
||||
rsa_password=encode_params.get("key_password"),
|
||||
embed_mode=encode_params["embed_mode"],
|
||||
dct_output_format=encode_params.get("dct_output_format", "png"),
|
||||
dct_color_mode=encode_params.get("dct_color_mode", "color"),
|
||||
channel_key=encode_params.get("channel_key"),
|
||||
progress_file=progress_file,
|
||||
)
|
||||
else:
|
||||
encode_result = subprocess_stego.encode(
|
||||
carrier_data=encode_params["carrier_data"],
|
||||
reference_data=encode_params["ref_data"],
|
||||
message=encode_params["message"],
|
||||
passphrase=encode_params["passphrase"],
|
||||
pin=encode_params.get("pin"),
|
||||
rsa_key_data=encode_params.get("rsa_key_data"),
|
||||
rsa_password=encode_params.get("key_password"),
|
||||
embed_mode=encode_params["embed_mode"],
|
||||
dct_output_format=encode_params.get("dct_output_format", "png"),
|
||||
dct_color_mode=encode_params.get("dct_color_mode", "color"),
|
||||
channel_key=encode_params.get("channel_key"),
|
||||
progress_file=progress_file,
|
||||
)
|
||||
|
||||
if not encode_result.success:
|
||||
_store_job(
|
||||
job_id,
|
||||
{
|
||||
"status": "error",
|
||||
"error": encode_result.error or "Encoding failed",
|
||||
"created": time.time(),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
# Determine output format
|
||||
embed_mode = encode_params["embed_mode"]
|
||||
dct_output_format = encode_params.get("dct_output_format", "png")
|
||||
dct_color_mode = encode_params.get("dct_color_mode", "color")
|
||||
|
||||
if embed_mode == "dct" and dct_output_format == "jpeg":
|
||||
output_ext = ".jpg"
|
||||
output_mime = "image/jpeg"
|
||||
else:
|
||||
output_ext = ".png"
|
||||
output_mime = "image/png"
|
||||
|
||||
filename = encode_result.filename
|
||||
if not filename:
|
||||
filename = generate_filename("stego", output_ext)
|
||||
elif embed_mode == "dct" and dct_output_format == "jpeg" and filename.endswith(".png"):
|
||||
filename = filename[:-4] + ".jpg"
|
||||
|
||||
# Store result
|
||||
file_id = secrets.token_urlsafe(16)
|
||||
TEMP_FILES[file_id] = {
|
||||
"data": encode_result.stego_data,
|
||||
"filename": filename,
|
||||
"timestamp": time.time(),
|
||||
"embed_mode": embed_mode,
|
||||
"output_format": dct_output_format if embed_mode == "dct" else "png",
|
||||
"color_mode": dct_color_mode if embed_mode == "dct" else None,
|
||||
"mime_type": output_mime,
|
||||
"channel_mode": encode_result.channel_mode,
|
||||
"channel_fingerprint": encode_result.channel_fingerprint,
|
||||
}
|
||||
|
||||
_store_job(
|
||||
job_id,
|
||||
{
|
||||
"status": "complete",
|
||||
"file_id": file_id,
|
||||
"created": time.time(),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
_store_job(
|
||||
job_id,
|
||||
{
|
||||
"status": "error",
|
||||
"error": str(e),
|
||||
"created": time.time(),
|
||||
},
|
||||
)
|
||||
finally:
|
||||
cleanup_progress_file(job_id)
|
||||
|
||||
|
||||
@app.route("/encode", methods=["GET", "POST"])
|
||||
@login_required
|
||||
def encode_page():
|
||||
if request.method == "POST":
|
||||
# Check if async mode requested
|
||||
is_async = request.form.get("async") == "true" or request.headers.get("X-Async") == "true"
|
||||
|
||||
try:
|
||||
# Get files
|
||||
ref_photo = request.files.get("reference_photo")
|
||||
@@ -956,7 +1109,9 @@ def encode_page():
|
||||
# Pre-check payload capacity BEFORE encode (fail fast)
|
||||
from stegasoo.steganography import will_fit_by_mode
|
||||
|
||||
payload_size = len(payload.data) if hasattr(payload, "data") else len(payload.encode("utf-8"))
|
||||
payload_size = (
|
||||
len(payload.data) if hasattr(payload, "data") else len(payload.encode("utf-8"))
|
||||
)
|
||||
fit_check = will_fit_by_mode(payload_size, carrier_data, embed_mode=embed_mode)
|
||||
if not fit_check.get("fits", True):
|
||||
error_msg = (
|
||||
@@ -972,8 +1127,35 @@ def encode_page():
|
||||
flash(error_msg, "error")
|
||||
return render_template("encode.html", has_qrcode_read=HAS_QRCODE_READ)
|
||||
|
||||
# v4.0.0: Include channel_key parameter
|
||||
# Use subprocess-isolated encode to prevent crashes
|
||||
# Build encode params for either sync or async
|
||||
encode_params = {
|
||||
"carrier_data": carrier_data,
|
||||
"ref_data": ref_data,
|
||||
"passphrase": passphrase,
|
||||
"pin": pin if pin else None,
|
||||
"rsa_key_data": rsa_key_data,
|
||||
"key_password": key_password,
|
||||
"embed_mode": embed_mode,
|
||||
"dct_output_format": dct_output_format if embed_mode == "dct" else "png",
|
||||
"dct_color_mode": dct_color_mode if embed_mode == "dct" else "color",
|
||||
"channel_key": channel_key,
|
||||
}
|
||||
|
||||
if payload_type == "file" and payload_file and payload_file.filename:
|
||||
encode_params["file_data"] = payload.data
|
||||
encode_params["file_name"] = payload.filename
|
||||
encode_params["file_mime"] = payload.mime_type
|
||||
else:
|
||||
encode_params["message"] = payload
|
||||
|
||||
# ASYNC MODE: Start background job and return JSON
|
||||
if is_async:
|
||||
job_id = generate_job_id()
|
||||
_store_job(job_id, {"status": "pending", "created": time.time()})
|
||||
_executor.submit(_run_encode_job, job_id, encode_params)
|
||||
return jsonify({"job_id": job_id, "status": "pending"})
|
||||
|
||||
# SYNC MODE: Run inline (original behavior)
|
||||
if payload_type == "file" and payload_file and payload_file.filename:
|
||||
encode_result = subprocess_stego.encode(
|
||||
carrier_data=carrier_data,
|
||||
@@ -988,7 +1170,7 @@ def encode_page():
|
||||
embed_mode=embed_mode,
|
||||
dct_output_format=dct_output_format if embed_mode == "dct" else "png",
|
||||
dct_color_mode=dct_color_mode if embed_mode == "dct" else "color",
|
||||
channel_key=channel_key, # v4.0.0
|
||||
channel_key=channel_key,
|
||||
)
|
||||
else:
|
||||
encode_result = subprocess_stego.encode(
|
||||
@@ -1002,7 +1184,7 @@ def encode_page():
|
||||
embed_mode=embed_mode,
|
||||
dct_output_format=dct_output_format if embed_mode == "dct" else "png",
|
||||
dct_color_mode=dct_color_mode if embed_mode == "dct" else "color",
|
||||
channel_key=channel_key, # v4.0.0
|
||||
channel_key=channel_key,
|
||||
)
|
||||
|
||||
# Check for subprocess errors
|
||||
@@ -1058,6 +1240,53 @@ def encode_page():
|
||||
return render_template("encode.html", has_qrcode_read=HAS_QRCODE_READ)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# ENCODE PROGRESS ENDPOINTS (v4.1.2)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@app.route("/encode/status/<job_id>")
|
||||
@login_required
|
||||
def encode_status(job_id):
|
||||
"""Get the status of an async encode job."""
|
||||
job = _get_job(job_id)
|
||||
if not job:
|
||||
return jsonify({"error": "Job not found"}), 404
|
||||
|
||||
response = {"status": job.get("status", "unknown")}
|
||||
|
||||
if job["status"] == "complete":
|
||||
response["file_id"] = job.get("file_id")
|
||||
elif job["status"] == "error":
|
||||
response["error"] = job.get("error", "Unknown error")
|
||||
|
||||
return jsonify(response)
|
||||
|
||||
|
||||
@app.route("/encode/progress/<job_id>")
|
||||
@login_required
|
||||
def encode_progress(job_id):
|
||||
"""Get the progress of an async encode job."""
|
||||
progress = read_progress(job_id)
|
||||
if progress:
|
||||
return jsonify(progress)
|
||||
|
||||
# No progress file yet - check job status
|
||||
job = _get_job(job_id)
|
||||
if not job:
|
||||
return jsonify({"error": "Job not found"}), 404
|
||||
|
||||
if job["status"] == "complete":
|
||||
return jsonify({"percent": 100, "phase": "complete"})
|
||||
elif job["status"] == "error":
|
||||
return jsonify({"percent": 0, "phase": "error", "error": job.get("error")})
|
||||
elif job["status"] == "pending":
|
||||
return jsonify({"percent": 0, "phase": "starting"})
|
||||
|
||||
# Running but no progress file yet
|
||||
return jsonify({"percent": 0, "phase": "initializing"})
|
||||
|
||||
|
||||
@app.route("/encode/result/<file_id>")
|
||||
@login_required
|
||||
def encode_result(file_id):
|
||||
@@ -1402,12 +1631,7 @@ def api_tools_strip_metadata():
|
||||
buffer = io.BytesIO(clean_data)
|
||||
filename = image_file.filename.rsplit(".", 1)[0] + "_clean.png"
|
||||
|
||||
return send_file(
|
||||
buffer,
|
||||
mimetype="image/png",
|
||||
as_attachment=True,
|
||||
download_name=filename
|
||||
)
|
||||
return send_file(buffer, mimetype="image/png", as_attachment=True, download_name=filename)
|
||||
except Exception as e:
|
||||
return jsonify({"success": False, "error": str(e)}), 400
|
||||
|
||||
@@ -1429,13 +1653,15 @@ def api_tools_exif():
|
||||
# Check if it's a JPEG (editable) or not
|
||||
is_jpeg = image_data[:2] == b"\xff\xd8"
|
||||
|
||||
return jsonify({
|
||||
"success": True,
|
||||
"filename": image_file.filename,
|
||||
"exif": exif,
|
||||
"editable": is_jpeg,
|
||||
"field_count": len(exif),
|
||||
})
|
||||
return jsonify(
|
||||
{
|
||||
"success": True,
|
||||
"filename": image_file.filename,
|
||||
"exif": exif,
|
||||
"editable": is_jpeg,
|
||||
"field_count": len(exif),
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
return jsonify({"success": False, "error": str(e)}), 400
|
||||
|
||||
@@ -1454,6 +1680,7 @@ def api_tools_exif_update():
|
||||
updates_json = request.form.get("updates", "{}")
|
||||
try:
|
||||
import json
|
||||
|
||||
updates = json.loads(updates_json)
|
||||
except json.JSONDecodeError:
|
||||
return jsonify({"success": False, "error": "Invalid updates JSON"}), 400
|
||||
@@ -1499,11 +1726,19 @@ def api_tools_exif_clear():
|
||||
clean_data = strip_image_metadata(image_data, output_format=output_format)
|
||||
|
||||
# Determine extension and mimetype
|
||||
ext_map = {"PNG": ("png", "image/png"), "JPEG": ("jpg", "image/jpeg"), "BMP": ("bmp", "image/bmp")}
|
||||
ext_map = {
|
||||
"PNG": ("png", "image/png"),
|
||||
"JPEG": ("jpg", "image/jpeg"),
|
||||
"BMP": ("bmp", "image/bmp"),
|
||||
}
|
||||
ext, mimetype = ext_map.get(output_format, ("png", "image/png"))
|
||||
|
||||
# Return as downloadable file
|
||||
stem = image_file.filename.rsplit(".", 1)[0] if "." in image_file.filename else image_file.filename
|
||||
stem = (
|
||||
image_file.filename.rsplit(".", 1)[0]
|
||||
if "." in image_file.filename
|
||||
else image_file.filename
|
||||
)
|
||||
buffer = io.BytesIO(clean_data)
|
||||
return send_file(
|
||||
buffer,
|
||||
@@ -1644,9 +1879,10 @@ def setup():
|
||||
@login_required
|
||||
def setup_recovery():
|
||||
"""Recovery key setup page (Step 2 of initial setup)."""
|
||||
from stegasoo.recovery import generate_recovery_key, hash_recovery_key, generate_recovery_qr
|
||||
import base64
|
||||
|
||||
from stegasoo.recovery import generate_recovery_key, generate_recovery_qr, hash_recovery_key
|
||||
|
||||
# Only allow during initial setup (no recovery key yet, first admin)
|
||||
if has_recovery_key():
|
||||
return redirect(url_for("index"))
|
||||
@@ -1724,9 +1960,10 @@ def recover():
|
||||
@admin_required
|
||||
def regenerate_recovery():
|
||||
"""Generate a new recovery key (replaces existing one)."""
|
||||
from stegasoo.recovery import generate_recovery_key, hash_recovery_key, generate_recovery_qr
|
||||
import base64
|
||||
|
||||
from stegasoo.recovery import generate_recovery_key, generate_recovery_qr, hash_recovery_key
|
||||
|
||||
if request.method == "POST":
|
||||
action = request.form.get("action")
|
||||
|
||||
@@ -1925,21 +2162,23 @@ def api_channel_keys():
|
||||
"""Get saved channel keys for current user (JSON API)."""
|
||||
current_user = get_current_user()
|
||||
keys = get_user_channel_keys(current_user.id)
|
||||
return jsonify({
|
||||
"success": True,
|
||||
"keys": [
|
||||
{
|
||||
"id": k.id,
|
||||
"name": k.name,
|
||||
"fingerprint": f"{k.channel_key[:4]}...{k.channel_key[-4:]}",
|
||||
"channel_key": k.channel_key,
|
||||
"last_used_at": k.last_used_at,
|
||||
}
|
||||
for k in keys
|
||||
],
|
||||
"can_save": can_save_channel_key(current_user.id),
|
||||
"max_keys": MAX_CHANNEL_KEYS,
|
||||
})
|
||||
return jsonify(
|
||||
{
|
||||
"success": True,
|
||||
"keys": [
|
||||
{
|
||||
"id": k.id,
|
||||
"name": k.name,
|
||||
"fingerprint": f"{k.channel_key[:4]}...{k.channel_key[-4:]}",
|
||||
"channel_key": k.channel_key,
|
||||
"last_used_at": k.last_used_at,
|
||||
}
|
||||
for k in keys
|
||||
],
|
||||
"can_save": can_save_channel_key(current_user.id),
|
||||
"max_keys": MAX_CHANNEL_KEYS,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@app.route("/api/channel/keys/<int:key_id>/use", methods=["POST"])
|
||||
|
||||
@@ -916,10 +916,184 @@ const Stegasoo = {
|
||||
});
|
||||
},
|
||||
|
||||
// ========================================================================
|
||||
// ASYNC ENCODE WITH PROGRESS (v4.1.2)
|
||||
// ========================================================================
|
||||
|
||||
/**
|
||||
* Submit encode form asynchronously with progress tracking
|
||||
* @param {HTMLFormElement} form - The encode form
|
||||
* @param {HTMLElement} btn - The submit button
|
||||
*/
|
||||
async submitEncodeAsync(form, btn) {
|
||||
const formData = new FormData(form);
|
||||
formData.append('async', 'true');
|
||||
|
||||
// Show progress modal
|
||||
this.showProgressModal('Encoding');
|
||||
|
||||
try {
|
||||
// Start encode job
|
||||
const response = await fetch('/encode', {
|
||||
method: 'POST',
|
||||
body: formData,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error('Failed to start encode');
|
||||
}
|
||||
|
||||
const result = await response.json();
|
||||
|
||||
if (result.error) {
|
||||
throw new Error(result.error);
|
||||
}
|
||||
|
||||
const jobId = result.job_id;
|
||||
|
||||
// Poll for progress
|
||||
await this.pollEncodeProgress(jobId);
|
||||
|
||||
} catch (error) {
|
||||
this.hideProgressModal();
|
||||
alert('Encode failed: ' + error.message);
|
||||
btn.disabled = false;
|
||||
btn.innerHTML = '<i class="bi bi-lock-fill me-2"></i>Encode';
|
||||
}
|
||||
},
|
||||
|
||||
/**
|
||||
* Poll encode progress until complete
|
||||
* @param {string} jobId - The job ID
|
||||
*/
|
||||
async pollEncodeProgress(jobId) {
|
||||
const progressBar = document.getElementById('progressBar');
|
||||
const progressText = document.getElementById('progressText');
|
||||
const phaseText = document.getElementById('progressPhase');
|
||||
|
||||
const poll = async () => {
|
||||
try {
|
||||
// Check status first
|
||||
const statusResponse = await fetch(`/encode/status/${jobId}`);
|
||||
const statusData = await statusResponse.json();
|
||||
|
||||
if (statusData.status === 'complete') {
|
||||
// Done - redirect to result
|
||||
this.updateProgress(100, 'Complete!');
|
||||
setTimeout(() => {
|
||||
window.location.href = `/encode/result/${statusData.file_id}`;
|
||||
}, 500);
|
||||
return;
|
||||
}
|
||||
|
||||
if (statusData.status === 'error') {
|
||||
throw new Error(statusData.error || 'Encode failed');
|
||||
}
|
||||
|
||||
// Get progress
|
||||
const progressResponse = await fetch(`/encode/progress/${jobId}`);
|
||||
const progressData = await progressResponse.json();
|
||||
|
||||
const percent = progressData.percent || 0;
|
||||
const phase = progressData.phase || 'processing';
|
||||
|
||||
this.updateProgress(percent, this.formatPhase(phase));
|
||||
|
||||
// Continue polling
|
||||
setTimeout(poll, 500);
|
||||
|
||||
} catch (error) {
|
||||
this.hideProgressModal();
|
||||
alert('Encode failed: ' + error.message);
|
||||
}
|
||||
};
|
||||
|
||||
await poll();
|
||||
},
|
||||
|
||||
/**
|
||||
* Format phase name for display
|
||||
*/
|
||||
formatPhase(phase) {
|
||||
const phases = {
|
||||
'starting': 'Starting...',
|
||||
'initializing': 'Initializing...',
|
||||
'embedding': 'Embedding data...',
|
||||
'saving': 'Saving image...',
|
||||
'finalizing': 'Finalizing...',
|
||||
'complete': 'Complete!',
|
||||
};
|
||||
return phases[phase] || phase;
|
||||
},
|
||||
|
||||
/**
|
||||
* Show progress modal
|
||||
*/
|
||||
showProgressModal(operation = 'Processing') {
|
||||
// Create modal if doesn't exist
|
||||
let modal = document.getElementById('progressModal');
|
||||
if (!modal) {
|
||||
modal = document.createElement('div');
|
||||
modal.id = 'progressModal';
|
||||
modal.className = 'modal fade';
|
||||
modal.setAttribute('data-bs-backdrop', 'static');
|
||||
modal.setAttribute('data-bs-keyboard', 'false');
|
||||
modal.innerHTML = `
|
||||
<div class="modal-dialog modal-dialog-centered">
|
||||
<div class="modal-content bg-dark text-light">
|
||||
<div class="modal-body p-4">
|
||||
<h5 class="mb-3" id="progressTitle">${operation}...</h5>
|
||||
<div class="progress mb-2" style="height: 24px;">
|
||||
<div id="progressBar" class="progress-bar progress-bar-striped progress-bar-animated bg-success"
|
||||
role="progressbar" style="width: 0%"></div>
|
||||
</div>
|
||||
<div class="d-flex justify-content-between text-muted small">
|
||||
<span id="progressPhase">Initializing...</span>
|
||||
<span id="progressText">0%</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
document.body.appendChild(modal);
|
||||
}
|
||||
|
||||
// Reset progress
|
||||
this.updateProgress(0, 'Initializing...');
|
||||
|
||||
// Show modal
|
||||
const bsModal = new bootstrap.Modal(modal);
|
||||
bsModal.show();
|
||||
},
|
||||
|
||||
/**
|
||||
* Hide progress modal
|
||||
*/
|
||||
hideProgressModal() {
|
||||
const modal = document.getElementById('progressModal');
|
||||
if (modal) {
|
||||
const bsModal = bootstrap.Modal.getInstance(modal);
|
||||
bsModal?.hide();
|
||||
}
|
||||
},
|
||||
|
||||
/**
|
||||
* Update progress bar and text
|
||||
*/
|
||||
updateProgress(percent, phase) {
|
||||
const progressBar = document.getElementById('progressBar');
|
||||
const progressText = document.getElementById('progressText');
|
||||
const phaseText = document.getElementById('progressPhase');
|
||||
|
||||
if (progressBar) progressBar.style.width = percent + '%';
|
||||
if (progressText) progressText.textContent = Math.round(percent) + '%';
|
||||
if (phaseText) phaseText.textContent = phase;
|
||||
},
|
||||
|
||||
// ========================================================================
|
||||
// INITIALIZATION HELPERS
|
||||
// ========================================================================
|
||||
|
||||
|
||||
initEncodePage() {
|
||||
this.initPasswordToggles();
|
||||
this.initRsaMethodToggle();
|
||||
@@ -937,27 +1111,23 @@ const Stegasoo = {
|
||||
generateBtnId: 'channelKeyGenerate'
|
||||
});
|
||||
|
||||
// Form submission with channel key validation
|
||||
// Form submission with async progress tracking (v4.1.2)
|
||||
const form = document.getElementById('encodeForm');
|
||||
const btn = document.getElementById('encodeBtn');
|
||||
form?.addEventListener('submit', (e) => {
|
||||
e.preventDefault();
|
||||
|
||||
if (!this.validateChannelKeyOnSubmit(form, 'channelSelect', 'channelKeyInput')) {
|
||||
e.preventDefault();
|
||||
return false;
|
||||
}
|
||||
|
||||
if (btn) {
|
||||
btn.disabled = true;
|
||||
const startTime = Date.now();
|
||||
const updateTimer = () => {
|
||||
const elapsed = Math.floor((Date.now() - startTime) / 1000);
|
||||
const mins = Math.floor(elapsed / 60);
|
||||
const secs = elapsed % 60;
|
||||
const timeStr = mins > 0 ? `${mins}:${secs.toString().padStart(2, '0')}` : `${secs}s`;
|
||||
btn.innerHTML = `<span class="spinner-border spinner-border-sm me-2"></span>Encoding... ${timeStr}`;
|
||||
};
|
||||
updateTimer();
|
||||
setInterval(updateTimer, 1000);
|
||||
btn.innerHTML = '<span class="spinner-border spinner-border-sm me-2"></span>Starting...';
|
||||
}
|
||||
|
||||
// Use async submission with progress tracking
|
||||
this.submitEncodeAsync(form, btn);
|
||||
});
|
||||
},
|
||||
|
||||
|
||||
@@ -111,6 +111,7 @@ def encode_operation(params: dict) -> dict:
|
||||
dct_output_format=params.get("dct_output_format", "png"),
|
||||
dct_color_mode=params.get("dct_color_mode", "color"),
|
||||
channel_key=resolved_channel_key, # v4.0.0
|
||||
progress_file=params.get("progress_file"), # v4.1.2
|
||||
)
|
||||
|
||||
# Build stats dict if available
|
||||
|
||||
@@ -47,6 +47,8 @@ import base64
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -233,6 +235,8 @@ class SubprocessStego:
|
||||
# Channel key (v4.0.0)
|
||||
channel_key: str | None = "auto",
|
||||
timeout: int | None = None,
|
||||
# Progress file (v4.1.2)
|
||||
progress_file: str | None = None,
|
||||
) -> EncodeResult:
|
||||
"""
|
||||
Encode a message or file into an image.
|
||||
@@ -268,6 +272,7 @@ class SubprocessStego:
|
||||
"dct_output_format": dct_output_format,
|
||||
"dct_color_mode": dct_color_mode,
|
||||
"channel_key": channel_key, # v4.0.0
|
||||
"progress_file": progress_file, # v4.1.2
|
||||
}
|
||||
|
||||
if file_data:
|
||||
@@ -496,3 +501,42 @@ def get_subprocess_stego() -> SubprocessStego:
|
||||
if _default_stego is None:
|
||||
_default_stego = SubprocessStego()
|
||||
return _default_stego
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Progress File Utilities (v4.1.2)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def generate_job_id() -> str:
|
||||
"""Generate a unique job ID for tracking encode/decode operations."""
|
||||
return str(uuid.uuid4())[:8]
|
||||
|
||||
|
||||
def get_progress_file_path(job_id: str) -> str:
|
||||
"""Get the progress file path for a job ID."""
|
||||
return str(Path(tempfile.gettempdir()) / f"stegasoo_progress_{job_id}.json")
|
||||
|
||||
|
||||
def read_progress(job_id: str) -> dict | None:
|
||||
"""
|
||||
Read progress from file for a job ID.
|
||||
|
||||
Returns:
|
||||
Progress dict with current, total, percent, phase, or None if not found
|
||||
"""
|
||||
progress_file = get_progress_file_path(job_id)
|
||||
try:
|
||||
with open(progress_file) as f:
|
||||
return json.load(f)
|
||||
except (FileNotFoundError, json.JSONDecodeError):
|
||||
return None
|
||||
|
||||
|
||||
def cleanup_progress_file(job_id: str) -> None:
|
||||
"""Remove progress file for a completed job."""
|
||||
progress_file = get_progress_file_path(job_id)
|
||||
try:
|
||||
Path(progress_file).unlink(missing_ok=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -54,6 +54,7 @@ cli = [
|
||||
"click>=8.0.0",
|
||||
"qrcode>=7.30",
|
||||
"piexif>=1.1.0",
|
||||
"rich>=13.0.0",
|
||||
]
|
||||
compression = [
|
||||
"lz4>=4.0.0",
|
||||
|
||||
@@ -390,7 +390,7 @@ gum style \
|
||||
--border-foreground 82 \
|
||||
--padding "0 2" \
|
||||
--align center \
|
||||
" ___ _____ ___ ___ _ ___ ___ ___" \
|
||||
" ___ _____ ___ ___ _ ___ ___ ___" \
|
||||
" / __||_ _|| __| / __| /_\\ / __| / _ \\ / _ \\" \
|
||||
" \\__ \\ | | | _| | (_ | / _ \\ \\__ \\ | (_) || (_) |" \
|
||||
" |___/ |_| |___| \\___//_/ \\_\\|___/ \\___/ \\___/" \
|
||||
|
||||
11
rpi/setup.sh
11
rpi/setup.sh
@@ -335,10 +335,13 @@ if systemctl is-active --quiet stegasoo 2>/dev/null; then
|
||||
STEGASOO_URL="http://$PI_IP:5000"
|
||||
fi
|
||||
echo ""
|
||||
echo -e "\033[0;36m ___ _____ ___ ___ _ ___ ___ ___\033[0m"
|
||||
echo -e "\033[0;36m / __||_ _|| __| / __| /_\\ / __| / _ \\ / _ \\\\\033[0m"
|
||||
echo -e "\033[0;36m \\__ \\ | | | _| | (_ | / _ \\ \\__ \\ | (_) || (_) |\033[0m"
|
||||
echo -e "\033[0;36m |___/ |_| |___| \\___//_/ \\_\\|___/ \\___/ \\___/\033[0m"
|
||||
echo -e "\033[0;90m . * . . * . * . * . * .\033[0m"
|
||||
echo -e "\033[0;36m ___ _____ ___ ___ _ ___ ___ ___\033[0m"
|
||||
echo -e "\033[0;36m / __||_ _|| __| / __| /_\\ / __| / _ \\ / _ \\\\\033[0m"
|
||||
echo -e "\033[0;36m \\__ \\ | | | _| | (_ | / _ \\ \\__ \\ | (_) || (_) |\033[0m"
|
||||
echo -e "\033[0;36m |___/ |_| |___| \\___//_/ \\_\\|___/ \\___/ \\___/\033[0m"
|
||||
echo ""
|
||||
echo -e "\033[0;90m * . * . * . * . * . *\033[0m"
|
||||
echo ""
|
||||
echo -e " \033[0;32m●\033[0m Stegasoo is running"
|
||||
echo -e " \033[0;33m$STEGASOO_URL\033[0m"
|
||||
|
||||
@@ -55,7 +55,32 @@ except ImportError:
|
||||
jio = None
|
||||
|
||||
# Import custom exceptions
|
||||
from .exceptions import InvalidMagicBytesError, ReedSolomonError as StegasooRSError
|
||||
from .exceptions import InvalidMagicBytesError
|
||||
from .exceptions import ReedSolomonError as StegasooRSError
|
||||
|
||||
# Progress reporting interval (write every N blocks)
|
||||
PROGRESS_INTERVAL = 50
|
||||
|
||||
|
||||
def _write_progress(progress_file: str | None, current: int, total: int, phase: str = "embedding"):
|
||||
"""Write progress to file for frontend polling."""
|
||||
if progress_file is None:
|
||||
return
|
||||
try:
|
||||
import json
|
||||
|
||||
with open(progress_file, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"current": current,
|
||||
"total": total,
|
||||
"percent": round((current / total) * 100, 1) if total > 0 else 0,
|
||||
"phase": phase,
|
||||
},
|
||||
f,
|
||||
)
|
||||
except Exception:
|
||||
pass # Don't let progress writing break encoding
|
||||
|
||||
|
||||
# ============================================================================
|
||||
@@ -189,7 +214,7 @@ def has_jpegio_support() -> bool:
|
||||
|
||||
# Check for reedsolo availability
|
||||
try:
|
||||
from reedsolo import RSCodec, ReedSolomonError
|
||||
from reedsolo import ReedSolomonError, RSCodec
|
||||
|
||||
HAS_REEDSOLO = True
|
||||
except ImportError:
|
||||
@@ -559,6 +584,7 @@ def embed_in_dct(
|
||||
seed: bytes,
|
||||
output_format: str = OUTPUT_FORMAT_PNG,
|
||||
color_mode: str = "color",
|
||||
progress_file: str | None = None,
|
||||
) -> tuple[bytes, DCTEmbedStats]:
|
||||
"""Embed data using DCT coefficient modification."""
|
||||
if output_format not in (OUTPUT_FORMAT_PNG, OUTPUT_FORMAT_JPEG):
|
||||
@@ -568,10 +594,12 @@ def embed_in_dct(
|
||||
color_mode = "color"
|
||||
|
||||
if output_format == OUTPUT_FORMAT_JPEG and HAS_JPEGIO:
|
||||
return _embed_jpegio(data, carrier_image, seed, color_mode)
|
||||
return _embed_jpegio(data, carrier_image, seed, color_mode, progress_file)
|
||||
|
||||
_check_scipy()
|
||||
return _embed_scipy_dct_safe(data, carrier_image, seed, output_format, color_mode)
|
||||
return _embed_scipy_dct_safe(
|
||||
data, carrier_image, seed, output_format, color_mode, progress_file
|
||||
)
|
||||
|
||||
|
||||
def _embed_scipy_dct_safe(
|
||||
@@ -580,6 +608,7 @@ def _embed_scipy_dct_safe(
|
||||
seed: bytes,
|
||||
output_format: str,
|
||||
color_mode: str = "color",
|
||||
progress_file: str | None = None,
|
||||
) -> tuple[bytes, DCTEmbedStats]:
|
||||
"""
|
||||
Embed using scipy DCT with safe memory handling.
|
||||
@@ -642,7 +671,7 @@ def _embed_scipy_dct_safe(
|
||||
gc.collect()
|
||||
|
||||
# Embed in Y channel
|
||||
Y_embedded = _embed_in_channel_safe(Y_padded, bits, block_order, blocks_x)
|
||||
Y_embedded = _embed_in_channel_safe(Y_padded, bits, block_order, blocks_x, progress_file)
|
||||
del Y_padded
|
||||
gc.collect()
|
||||
|
||||
@@ -666,7 +695,7 @@ def _embed_scipy_dct_safe(
|
||||
del image
|
||||
gc.collect()
|
||||
|
||||
embedded = _embed_in_channel_safe(padded, bits, block_order, blocks_x)
|
||||
embedded = _embed_in_channel_safe(padded, bits, block_order, blocks_x, progress_file)
|
||||
del padded
|
||||
gc.collect()
|
||||
|
||||
@@ -699,6 +728,7 @@ def _embed_in_channel_safe(
|
||||
bits: list,
|
||||
block_order: list,
|
||||
blocks_x: int,
|
||||
progress_file: str | None = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Embed bits in channel using safe DCT operations.
|
||||
@@ -711,8 +741,9 @@ def _embed_in_channel_safe(
|
||||
result = np.array(channel, dtype=np.float64, copy=True, order="C")
|
||||
|
||||
bit_idx = 0
|
||||
total_blocks = len(block_order)
|
||||
|
||||
for block_num in block_order:
|
||||
for block_idx, block_num in enumerate(block_order):
|
||||
if bit_idx >= len(bits):
|
||||
break
|
||||
|
||||
@@ -748,6 +779,14 @@ def _embed_in_channel_safe(
|
||||
# Clean up this iteration
|
||||
del block, dct_block, modified_block
|
||||
|
||||
# Report progress periodically
|
||||
if progress_file and block_idx % PROGRESS_INTERVAL == 0:
|
||||
_write_progress(progress_file, block_idx, total_blocks, "embedding")
|
||||
|
||||
# Final progress update
|
||||
if progress_file:
|
||||
_write_progress(progress_file, total_blocks, total_blocks, "finalizing")
|
||||
|
||||
# Force garbage collection
|
||||
gc.collect()
|
||||
|
||||
@@ -804,6 +843,7 @@ def _embed_jpegio(
|
||||
carrier_image: bytes,
|
||||
seed: bytes,
|
||||
color_mode: str = "color",
|
||||
progress_file: str | None = None,
|
||||
) -> tuple[bytes, DCTEmbedStats]:
|
||||
"""Embed using jpegio for proper JPEG coefficient modification."""
|
||||
import os
|
||||
@@ -861,6 +901,9 @@ def _embed_jpegio(
|
||||
)
|
||||
|
||||
coefs_used = 0
|
||||
total_bits = len(bits)
|
||||
progress_interval = max(total_bits // 20, 100) # Report ~20 times or every 100 bits
|
||||
|
||||
for bit_idx, pos_idx in enumerate(order):
|
||||
if bit_idx >= len(bits):
|
||||
break
|
||||
@@ -876,6 +919,14 @@ def _embed_jpegio(
|
||||
|
||||
coefs_used += 1
|
||||
|
||||
# Report progress periodically
|
||||
if progress_file and bit_idx % progress_interval == 0:
|
||||
_write_progress(progress_file, bit_idx, total_bits, "embedding")
|
||||
|
||||
# Final progress before save
|
||||
if progress_file:
|
||||
_write_progress(progress_file, total_bits, total_bits, "saving")
|
||||
|
||||
jio.write(jpeg, output_path)
|
||||
|
||||
with open(output_path, "rb") as f:
|
||||
@@ -971,8 +1022,8 @@ def _extract_scipy_dct_safe(stego_image: bytes, seed: bytes) -> bytes:
|
||||
total_needed = (HEADER_SIZE + data_length) * 8
|
||||
if len(all_bits) >= total_needed:
|
||||
break
|
||||
except ValueError:
|
||||
pass
|
||||
except (ValueError, InvalidMagicBytesError):
|
||||
pass # RS-protected format has length prefix first, not magic bytes
|
||||
|
||||
del padded
|
||||
gc.collect()
|
||||
@@ -997,6 +1048,7 @@ def _extract_scipy_dct_safe(stego_image: bytes, seed: bytes) -> bytes:
|
||||
|
||||
# Count occurrences of each unique copy
|
||||
from collections import Counter
|
||||
|
||||
counter = Counter(copies)
|
||||
best_header, count = counter.most_common(1)[0]
|
||||
|
||||
@@ -1009,9 +1061,13 @@ def _extract_scipy_dct_safe(stego_image: bytes, seed: bytes) -> bytes:
|
||||
|
||||
# Sanity check: both lengths should be reasonable
|
||||
max_reasonable = (len(all_bits) // 8) - RS_LENGTH_PREFIX_SIZE
|
||||
if (raw_payload_length > 0 and raw_payload_length <= max_reasonable and
|
||||
rs_encoded_length > 0 and rs_encoded_length <= max_reasonable and
|
||||
rs_encoded_length >= raw_payload_length):
|
||||
if (
|
||||
raw_payload_length > 0
|
||||
and raw_payload_length <= max_reasonable
|
||||
and rs_encoded_length > 0
|
||||
and rs_encoded_length <= max_reasonable
|
||||
and rs_encoded_length >= raw_payload_length
|
||||
):
|
||||
# This looks like RS-protected format
|
||||
total_bits_needed = (RS_LENGTH_PREFIX_SIZE + rs_encoded_length) * 8
|
||||
|
||||
@@ -1088,6 +1144,7 @@ def _extract_jpegio(stego_image: bytes, seed: bytes) -> bytes:
|
||||
|
||||
# Extract 3 copies and use majority voting
|
||||
from collections import Counter
|
||||
|
||||
copies = []
|
||||
for i in range(RS_LENGTH_COPIES):
|
||||
start = i * RS_LENGTH_HEADER_SIZE
|
||||
@@ -1104,9 +1161,13 @@ def _extract_jpegio(stego_image: bytes, seed: bytes) -> bytes:
|
||||
|
||||
# Sanity check
|
||||
max_reasonable = (len(all_positions) // 8) - RS_LENGTH_PREFIX_SIZE
|
||||
if (raw_payload_length > 0 and raw_payload_length <= max_reasonable and
|
||||
rs_encoded_length > 0 and rs_encoded_length <= max_reasonable and
|
||||
rs_encoded_length >= raw_payload_length):
|
||||
if (
|
||||
raw_payload_length > 0
|
||||
and raw_payload_length <= max_reasonable
|
||||
and rs_encoded_length > 0
|
||||
and rs_encoded_length <= max_reasonable
|
||||
and rs_encoded_length >= raw_payload_length
|
||||
):
|
||||
total_bits_needed = (RS_LENGTH_PREFIX_SIZE + rs_encoded_length) * 8
|
||||
|
||||
if len(all_positions) >= total_bits_needed:
|
||||
|
||||
@@ -37,6 +37,7 @@ def encode(
|
||||
dct_output_format: str = "png",
|
||||
dct_color_mode: str = "color",
|
||||
channel_key: str | bool | None = None,
|
||||
progress_file: str | None = None,
|
||||
) -> EncodeResult:
|
||||
"""
|
||||
Encode a message or file into an image.
|
||||
@@ -118,6 +119,7 @@ def encode(
|
||||
embed_mode=embed_mode,
|
||||
dct_output_format=dct_output_format,
|
||||
dct_color_mode=dct_color_mode,
|
||||
progress_file=progress_file,
|
||||
)
|
||||
|
||||
# Generate filename
|
||||
|
||||
@@ -39,6 +39,31 @@ from .debug import debug
|
||||
from .exceptions import CapacityError, EmbeddingError
|
||||
from .models import EmbedStats, FilePayload
|
||||
|
||||
# Progress reporting interval
|
||||
PROGRESS_INTERVAL = 1000 # Write every N pixels for LSB
|
||||
|
||||
|
||||
def _write_progress(progress_file: str | None, current: int, total: int, phase: str = "embedding"):
|
||||
"""Write progress to file for frontend polling."""
|
||||
if progress_file is None:
|
||||
return
|
||||
try:
|
||||
import json
|
||||
|
||||
with open(progress_file, "w") as f:
|
||||
json.dump(
|
||||
{
|
||||
"current": current,
|
||||
"total": total,
|
||||
"percent": round((current / total) * 100, 1) if total > 0 else 0,
|
||||
"phase": phase,
|
||||
},
|
||||
f,
|
||||
)
|
||||
except Exception:
|
||||
pass # Don't let progress writing break encoding
|
||||
|
||||
|
||||
# Lossless formats that preserve LSB data
|
||||
LOSSLESS_FORMATS = {"PNG", "BMP", "TIFF"}
|
||||
|
||||
@@ -526,6 +551,7 @@ def embed_in_image(
|
||||
embed_mode: str = EMBED_MODE_LSB,
|
||||
dct_output_format: str = DCT_OUTPUT_PNG,
|
||||
dct_color_mode: str = "color",
|
||||
progress_file: str | None = None,
|
||||
) -> tuple[bytes, Union[EmbedStats, "DCTEmbedStats"], str]:
|
||||
"""
|
||||
Embed data into an image using specified mode.
|
||||
@@ -579,6 +605,7 @@ def embed_in_image(
|
||||
pixel_key,
|
||||
output_format=dct_output_format,
|
||||
color_mode=dct_color_mode,
|
||||
progress_file=progress_file,
|
||||
)
|
||||
|
||||
# Determine extension based on output format
|
||||
@@ -594,7 +621,7 @@ def embed_in_image(
|
||||
return stego_bytes, dct_stats, ext
|
||||
|
||||
# LSB MODE
|
||||
return _embed_lsb(data, image_data, pixel_key, bits_per_channel, output_format)
|
||||
return _embed_lsb(data, image_data, pixel_key, bits_per_channel, output_format, progress_file)
|
||||
|
||||
|
||||
def _embed_lsb(
|
||||
@@ -603,6 +630,7 @@ def _embed_lsb(
|
||||
pixel_key: bytes,
|
||||
bits_per_channel: int = 1,
|
||||
output_format: str | None = None,
|
||||
progress_file: str | None = None,
|
||||
) -> tuple[bytes, EmbedStats, str]:
|
||||
"""
|
||||
Embed data using LSB steganography (internal implementation).
|
||||
@@ -659,8 +687,9 @@ def _embed_lsb(
|
||||
|
||||
bit_idx = 0
|
||||
modified_pixels = 0
|
||||
total_pixels_to_process = len(selected_indices)
|
||||
|
||||
for pixel_idx in selected_indices:
|
||||
for progress_idx, pixel_idx in enumerate(selected_indices):
|
||||
if bit_idx >= len(binary_data):
|
||||
break
|
||||
|
||||
@@ -690,6 +719,16 @@ def _embed_lsb(
|
||||
new_pixels[pixel_idx] = (r, g, b)
|
||||
modified_pixels += 1
|
||||
|
||||
# Report progress periodically
|
||||
if progress_file and progress_idx % PROGRESS_INTERVAL == 0:
|
||||
_write_progress(progress_file, progress_idx, total_pixels_to_process, "embedding")
|
||||
|
||||
# Final progress before save
|
||||
if progress_file:
|
||||
_write_progress(
|
||||
progress_file, total_pixels_to_process, total_pixels_to_process, "saving"
|
||||
)
|
||||
|
||||
debug.print(f"Modified {modified_pixels} pixels (out of {len(selected_indices)} selected)")
|
||||
|
||||
stego_img = Image.new("RGB", img.size)
|
||||
|
||||
Reference in New Issue
Block a user