ghostwire/whatsapp_suite/zip_export.py

"""Export all local message rows to ZIP via the native WebView JavaScript bridge.

No CDP, UI automation, app restart, or phone history sync. The ZIP contains
messages.jsonl and manifest.json; media binaries are not downloaded.
"""
import argparse
import datetime as dt
import json
import os
from pathlib import Path
import time
import zipfile

from .history import expression
from .live import data_chunk_expression, run_math


def evaluate(execute, value, session):
    """Run the fixed async database probe and stream its result across the bridge."""
    key = '__whatsapp_zip_' + session
    literal = json.dumps(key)
    start = '''(() => {
      const key=__KEY__,id=__ID__,state={probe_id:id,state:'pending'};window[key]=state;
      Promise.resolve(__VALUE__).then(value=>{
        state.json=JSON.stringify(value);state.state='ready';
      }).catch(error=>{state.state='error';state.error=String(error).slice(0,500);});
      return {probe_id:id,state:'started'};
    })()'''.replace('__KEY__', literal).replace('__ID__', json.dumps(session)).replace('__VALUE__', value)
    poll = '''(() => {const v=window[__KEY__];return {probe_id:v?.probe_id,state:v?.state,error:v?.error,length:v?.json?.length};})()'''.replace('__KEY__', literal)
    try:
        initial = json.loads(execute(start, 'export batch startup'))
        if initial.get('probe_id') != session or initial.get('state') != 'started':
            raise ValueError('Export startup response mismatch')
        deadline = time.monotonic() + 60
        while True:
            meta = json.loads(execute(poll, 'export batch poll'))
            if meta.get('probe_id') != session:
                raise ValueError('Export session mismatch')
            if meta.get('state') == 'error':
                raise RuntimeError(meta.get('error', 'JavaScript export failed'))
            if meta.get('state') == 'ready':
                break
            if time.monotonic() >= deadline:
                raise TimeoutError('JavaScript export batch timed out')
            time.sleep(.05)
        length = meta.get('length')
        if type(length) is not int or not 0 < length <= 16 * 1024 * 1024:
            raise ValueError('Export batch exceeds 16 MiB transport bound; reduce --batch-size')
        offset, parts = 0, []
        while offset < length:
            part = json.loads(execute(data_chunk_expression(key, offset), 'export batch chunk'))
            following = part.get('next')
            if (part.get('probe_id') != session or part.get('offset') != offset
                    or type(following) is not int or not offset < following <= length
                    or not isinstance(part.get('chunk'), str)
                    or len(part['chunk'].encode('utf-16-le', errors='surrogatepass')) // 2 != following-offset):
                raise ValueError('Export transport chunk mismatch')
            parts.append(part['chunk'])
            offset = following
        return json.loads(''.join(parts))
    finally:
        execute('delete window[' + literal + ']', 'export JavaScript state cleanup')


def write_messages(read, archive, batch, redact):
    """Stream a bounded keyset scan, preserving a record for every row error."""
    bounds = read({'mode': 'bounds'})
    count, upper = bounds['count'], bounds['upper']
    after, rows, errors, seen = None, 0, 0, set()
    with archive.open('messages.jsonl', 'w', force_zip64=True) as stream:
        while upper is not None and after != upper:
            page = read({'mode': 'batch', 'upper': upper, 'after': after, 'batch': batch, 'redact': redact})
            keys, messages = page['keys'], page['messages']
            if not keys:
                raise RuntimeError('History cursor ended before the initial upper bound')
            if len(keys) != len(messages) or page['last'] != keys[-1] or page['last'] == after:
                raise RuntimeError('Invalid or nonprogressing history page')
            for key, message in zip(keys, messages):
                if key in seen or message.get('id') != key:
                    raise RuntimeError('Duplicate key or mismatched history message')
                seen.add(key)
                stream.write((json.dumps(message, ensure_ascii=True) + '\n').encode('utf-8'))
                rows += 1
                errors += int('error' in message)
            after = page['last']
            print(json.dumps({'exported_rows': rows, 'initial_rows': count, 'row_errors': errors}), flush=True)
    final = read({'mode': 'bounds'})
    return {'rows': rows, 'row_errors': errors, 'initial_row_count': count,
            'final_row_count': final['count'], 'initial_upper_key': upper,
            'reached_initial_upper': after == upper, 'snapshot_atomic': False,
            'count_matches_initial': rows == count}


def export_zip(output, pid=None, batch=100, redact=False):
    output = Path(output).resolve()
    if not 1 <= batch <= 500:
        raise ValueError('Batch size must be 1..500')
    partial = output.with_suffix(output.suffix + '.partial')
    report_path = output.with_suffix(output.suffix + '.report.json')
    if output.exists() or partial.exists() or report_path.exists():
        raise ValueError('Output or partial/report exists; select a new --output')
    output.parent.mkdir(parents=True, exist_ok=True)
    summary = {'cdp': False, 'ui_automation': False, 'transport': 'native WebView ExecuteScriptAsync',
               'source': 'WhatsApp JavaScript message table getter and decoder',
               'scope': 'All local message rows up to initial upper primary key; no phone sync or media download',
               'created_utc': dt.datetime.now(dt.timezone.utc).isoformat(), 'text_redacted': redact}
    def operation(execute, session, bridge):
        with zipfile.ZipFile(partial, 'x', compression=zipfile.ZIP_DEFLATED, allowZip64=True) as archive:
            summary.update(write_messages(lambda options: evaluate(execute, expression(options), session), archive, batch, redact))
    bridge = run_math(pid, 'history-export', _operation=operation)
    summary['bridge'] = bridge
    summary['passed'] = bool(bridge['passed'] and summary.get('reached_initial_upper') and not summary.get('row_errors'))
    summary['status'] = 'completed' if summary['passed'] else 'completed_with_errors' if bridge['passed'] else 'failed_partial_export'
    try:
        if bridge['passed']:
            with zipfile.ZipFile(partial, 'a', compression=zipfile.ZIP_DEFLATED) as archive:
                archive.writestr('manifest.json', json.dumps(summary, indent=2))
            with zipfile.ZipFile(partial) as archive:
                bad = archive.testzip()
                if bad:
                    raise RuntimeError('ZIP integrity failed: ' + bad)
            # Publish without overwriting an output created during the live scan.
            os.link(partial, output)
            partial.unlink()
            summary['output'] = str(output)
        else:
            summary['partial_output'] = str(partial)
    except Exception as error:
        summary.update(passed=False, status='failed_partial_export', error=str(error), partial_output=str(partial))
    report_path.write_text(json.dumps(summary, indent=2) + '\n', encoding='utf-8')
    return summary


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('exports') / ('whatsapp-' + dt.datetime.now().strftime('%Y%m%d-%H%M%S') + '.zip'))
    parser.add_argument('--pid', type=int)
    parser.add_argument('--batch-size', type=int, default=100)
    parser.add_argument('--redact-text', action='store_true')
    args = parser.parse_args(argv)
    report = export_zip(args.output, args.pid, args.batch_size, args.redact_text)
    print(json.dumps({k: report[k] for k in ('passed', 'status', 'rows', 'row_errors', 'output', 'partial_output', 'error') if k in report}))
    return 0 if report['passed'] else 1