#!/usr/bin/env python3
"""
RTExit Auto-Documentation Engine
Automatically logs all activities, commands, and findings.
Usage:
  python3 autodoc_engine.py log --skill rt-osint --cmd "amass enum -d target.com" --output "..."
  python3 autodoc_engine.py timeline
  python3 autodoc_engine.py custody --finding F-001 --evidence "screenshot.png"
"""

import argparse
import datetime
import hashlib
import json
import os


OUTPUT_DIR = os.environ.get('RTEXIT_OUTPUT', '_rtexit-output')
DOCS_DIR = os.path.join(OUTPUT_DIR, 'docs')
TIMELINE_FILE = os.path.join(DOCS_DIR, 'engagement', 'timeline.md')
CUSTODY_FILE = os.path.join(DOCS_DIR, 'evidence', 'chain-of-custody.md')
TERMINAL_LOG_DIR = os.path.join(DOCS_DIR, 'evidence', 'terminal-logs')


def ensure_dirs():
    for d in [
        os.path.join(DOCS_DIR, 'engagement'),
        os.path.join(DOCS_DIR, 'reconnaissance'),
        os.path.join(DOCS_DIR, 'findings'),
        os.path.join(DOCS_DIR, 'evidence', 'screenshots'),
        os.path.join(DOCS_DIR, 'evidence', 'terminal-logs'),
        os.path.join(DOCS_DIR, 'evidence', 'http-logs'),
        os.path.join(DOCS_DIR, 'attack-chains'),
        os.path.join(DOCS_DIR, 'reports'),
    ]:
        os.makedirs(d, exist_ok=True)

    # Initialize timeline if not exists
    if not os.path.exists(TIMELINE_FILE):
        with open(TIMELINE_FILE, 'w', encoding='utf-8') as f:
            f.write("# Engagement Timeline\n\n")
            f.write("| Timestamp | Phase | Skill | Activity | Finding |\n")
            f.write("|-----------|-------|-------|----------|---------|\n")

    # Initialize custody log if not exists
    if not os.path.exists(CUSTODY_FILE):
        with open(CUSTODY_FILE, 'w', encoding='utf-8') as f:
            f.write("# Chain of Custody Log\n\n")
            f.write("| Timestamp | Finding | Evidence | SHA-256 Hash | Operator |\n")
            f.write("|-----------|---------|----------|--------------|----------|\n")


def sha256_file(path: str) -> str:
    if not os.path.exists(path):
        return "file-not-found"
    h = hashlib.sha256()
    with open(path, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            h.update(chunk)
    return h.hexdigest()


def sha256_str(s: str) -> str:
    return hashlib.sha256(s.encode()).hexdigest()


def cmd_log(args):
    ensure_dirs()
    ts = datetime.datetime.now().isoformat(timespec='seconds')
    date_str = datetime.date.today().strftime('%Y%m%d')

    # Log to timeline
    finding_ref = args.finding or '-'
    activity = args.cmd[:80] if args.cmd else args.note or 'activity'
    with open(TIMELINE_FILE, 'a', encoding='utf-8') as f:
        f.write(f"| {ts} | {args.phase or '-'} | {args.skill or '-'} | `{activity}` | {finding_ref} |\n")

    # Save terminal log
    if args.cmd:
        log_filename = f"{date_str}_{args.skill or 'unknown'}_{ts.replace(':', '-')}.txt"
        log_path = os.path.join(TERMINAL_LOG_DIR, log_filename)
        with open(log_path, 'w', encoding='utf-8') as f:
            f.write(f"# Terminal Log\n")
            f.write(f"Timestamp : {ts}\n")
            f.write(f"Skill     : {args.skill or 'unknown'}\n")
            f.write(f"Phase     : {args.phase or 'unknown'}\n")
            f.write(f"Command   : {args.cmd}\n")
            f.write(f"Finding   : {finding_ref}\n\n")
            f.write("## Output\n```\n")
            f.write(args.output or '(no output captured)')
            f.write("\n```\n")

        # Chain of custody
        evidence_hash = sha256_str(args.output or '')
        with open(CUSTODY_FILE, 'a', encoding='utf-8') as f:
            f.write(f"| {ts} | {finding_ref} | {log_filename} | `{evidence_hash[:16]}...` | {args.operator or '-'} |\n")

        print(f"✅ Logged: {log_filename}")
    else:
        print(f"✅ Timeline entry added: {ts}")


def cmd_timeline(args):
    ensure_dirs()
    if os.path.exists(TIMELINE_FILE):
        print(open(TIMELINE_FILE).read())
    else:
        print("No timeline yet.")


def cmd_custody(args):
    ensure_dirs()

    if args.evidence and os.path.exists(args.evidence):
        evidence_hash = sha256_file(args.evidence)
        evidence_name = os.path.basename(args.evidence)
    else:
        evidence_hash = sha256_str(args.evidence or '')
        evidence_name = args.evidence or 'unknown'

    ts = datetime.datetime.now().isoformat(timespec='seconds')
    with open(CUSTODY_FILE, 'a', encoding='utf-8') as f:
        f.write(f"| {ts} | {args.finding or '-'} | {evidence_name} | `{evidence_hash[:32]}` | {args.operator or '-'} |\n")

    print(f"✅ Evidence logged: {evidence_name}")
    print(f"   SHA-256: {evidence_hash}")


def cmd_init(args):
    """Initialize output directory for a new engagement."""
    ensure_dirs()

    # Create engagement files
    engagement_info = {
        'ref': args.ref or '',
        'client': args.client or '',
        'start': datetime.date.today().isoformat(),
        'methodology': args.methodology or 'ptes',
        'scope': args.scope or 'blackbox',
    }

    info_path = os.path.join(DOCS_DIR, 'engagement', 'engagement-info.json')
    with open(info_path, 'w') as f:
        json.dump(engagement_info, f, indent=2)

    scope_path = os.path.join(DOCS_DIR, 'engagement', 'scope.md')
    if not os.path.exists(scope_path):
        with open(scope_path, 'w') as f:
            f.write(f"# Scope — {args.ref or 'ENGAGEMENT'}\n\n")
            f.write("## In-Scope Targets\n\n| Target | Type | Priority |\n|--------|------|----------|\n| | | |\n\n")
            f.write("## Out-of-Scope\n\n-\n\n")
            f.write("## Notes\n\n")

    print(f"✅ Engagement initialized: {args.ref}")
    print(f"   Output directory: {os.path.abspath(OUTPUT_DIR)}")


def main():
    parser = argparse.ArgumentParser(description='RTExit Auto-Documentation Engine')
    sub = parser.add_subparsers(dest='command')

    # log
    p_log = sub.add_parser('log', help='Log an activity')
    p_log.add_argument('--skill', help='Skill name')
    p_log.add_argument('--phase', help='Current phase')
    p_log.add_argument('--cmd', help='Command executed')
    p_log.add_argument('--output', help='Command output')
    p_log.add_argument('--finding', help='Related finding ID')
    p_log.add_argument('--operator', help='Operator name')
    p_log.add_argument('--note', help='Free text note')

    # timeline
    sub.add_parser('timeline', help='Show engagement timeline')

    # custody
    p_cust = sub.add_parser('custody', help='Log evidence chain of custody')
    p_cust.add_argument('--finding', required=True, help='Finding ID')
    p_cust.add_argument('--evidence', required=True, help='Evidence file path or description')
    p_cust.add_argument('--operator', help='Operator name')

    # init
    p_init = sub.add_parser('init', help='Initialize new engagement')
    p_init.add_argument('--ref', help='Engagement reference')
    p_init.add_argument('--client', help='Client name')
    p_init.add_argument('--methodology', default='ptes')
    p_init.add_argument('--scope', default='blackbox')

    args = parser.parse_args()

    if args.command == 'log':
        cmd_log(args)
    elif args.command == 'timeline':
        cmd_timeline(args)
    elif args.command == 'custody':
        cmd_custody(args)
    elif args.command == 'init':
        cmd_init(args)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()
