"""OCPP 양방향 프록시 / 로깅 서버 (인증 포함)

충전기 ↔ [이 프록시] ↔ CPO 서버
양방향 WebSocket 메시지를 투명하게 중계하면서 모든 통신을 로깅.

실행:
  pip3 install websockets aiohttp
  python3 ocpp_proxy_server.py

구조:
  - 포트 9000: WebSocket 프록시 (충전기 연결, 인증 불필요)
  - 포트 9001: 관리 웹서버 (로그인 필수)
"""

import asyncio
import json
import os
import time
import struct
import logging
import hashlib
import hmac
import secrets
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Dict, Optional, Set
from dataclasses import dataclass, asdict

import websockets
from aiohttp import web

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  설정
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

PROXY_PORT = 9002
WEB_PORT = 9003
LOG_DIR = "ocpp_logs"
CONFIG_FILE = "proxy_config.json"
USERS_FILE = "proxy_users.json"
TOKEN_SECRET = secrets.token_hex(32)
TOKEN_EXPIRE_HOURS = 24

DEFAULT_CONFIG = {
    "target_url": "ws://cp.e-csp.co.kr/ocppext",
    "target_name": "e-CSP CPO",
    "log_format": "both",
    "log_enabled": True,
    "pcap_enabled": False,
    "ocpp_subprotocol": "ocpp1.6",
    "max_log_size_mb": 100,
    "my_ocpp_server": "ws://s1.byunc.com/steve/websocket/CentralSystemService",
    "my_dashboard": "http://s1.byunc.com/dashboard",
}

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")
logger = logging.getLogger("ocpp-proxy")
MSG_TYPES = {2: "CALL", 3: "CALLRESULT", 4: "CALLERROR"}
C="\033[96m";G="\033[92m";Y="\033[93m";R="\033[91m";M="\033[95m";DIM="\033[2m";E="\033[0m"


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  사용자 관리
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _hash_pw(password: str, salt: str = None) -> tuple:
    if not salt:
        salt = secrets.token_hex(16)
    hashed = hashlib.pbkdf2_hmac("sha256", password.encode(), salt.encode(), 100000).hex()
    return hashed, salt

def _verify_pw(password: str, hashed: str, salt: str) -> bool:
    check, _ = _hash_pw(password, salt)
    return hmac.compare_digest(check, hashed)

def _create_token(user_id: str, username: str, role: str) -> str:
    payload = {
        "uid": user_id, "user": username, "role": role,
        "exp": (datetime.now(timezone.utc) + timedelta(hours=TOKEN_EXPIRE_HOURS)).timestamp(),
        "nonce": secrets.token_hex(8),
    }
    data = json.dumps(payload, separators=(",", ":"))
    import base64
    b64 = base64.urlsafe_b64encode(data.encode()).decode()
    sig = hmac.new(TOKEN_SECRET.encode(), b64.encode(), hashlib.sha256).hexdigest()[:32]
    return f"{b64}.{sig}"

def _decode_token(token: str) -> Optional[dict]:
    try:
        import base64
        parts = token.split(".")
        if len(parts) != 2:
            return None
        b64, sig = parts
        expected = hmac.new(TOKEN_SECRET.encode(), b64.encode(), hashlib.sha256).hexdigest()[:32]
        if not hmac.compare_digest(sig, expected):
            return None
        payload = json.loads(base64.urlsafe_b64decode(b64))
        if payload.get("exp", 0) < datetime.now(timezone.utc).timestamp():
            return None
        return payload
    except Exception:
        return None


class UserManager:
    def __init__(self):
        self.users = {}
        self.load()
        if not self.users:
            self.add_user("admin", "admin1234", "관리자", "admin")
            logger.info(f"{G}초기 관리자 계정 생성: admin / admin1234{E}")

    def load(self):
        try:
            if os.path.exists(USERS_FILE):
                with open(USERS_FILE, "r") as f:
                    self.users = json.load(f)
        except Exception:
            self.users = {}

    def save(self):
        with open(USERS_FILE, "w") as f:
            json.dump(self.users, f, indent=2, ensure_ascii=False)

    def add_user(self, username, password, display_name="", role="viewer"):
        hashed, salt = _hash_pw(password)
        self.users[username] = {
            "username": username,
            "hashed": hashed,
            "salt": salt,
            "display_name": display_name or username,
            "role": role,
            "is_active": True,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "last_login": None,
        }
        self.save()

    def verify(self, username, password):
        user = self.users.get(username)
        if not user or not user.get("is_active"):
            return None
        if _verify_pw(password, user["hashed"], user["salt"]):
            user["last_login"] = datetime.now(timezone.utc).isoformat()
            self.save()
            return user
        return None

    def update_user(self, username, updates: dict):
        user = self.users.get(username)
        if not user:
            return False
        if "password" in updates:
            hashed, salt = _hash_pw(updates.pop("password"))
            user["hashed"] = hashed
            user["salt"] = salt
        for k in ("display_name", "role", "is_active"):
            if k in updates:
                user[k] = updates[k]
        self.save()
        return True

    def delete_user(self, username):
        if username in self.users:
            del self.users[username]
            self.save()
            return True
        return False

    def list_users(self):
        return [
            {k: v for k, v in u.items() if k not in ("hashed", "salt")}
            for u in self.users.values()
        ]


user_mgr = UserManager()


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  인증 미들웨어
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

PUBLIC_PATHS = {"/", "/api/auth/login"}

def _get_token(request):
    auth = request.headers.get("Authorization", "")
    if auth.startswith("Bearer "):
        return auth[7:]
    return request.cookies.get("token")

@web.middleware
async def auth_middleware(request, handler):
    if request.path in PUBLIC_PATHS or request.path.startswith("/ws/"):
        return await handler(request)
    token = _get_token(request)
    if not token:
        return web.json_response({"error": "인증 필요"}, status=401)
    payload = _decode_token(token)
    if not payload:
        return web.json_response({"error": "만료된 토큰"}, status=401)
    user = user_mgr.users.get(payload.get("user"))
    if not user or not user.get("is_active"):
        return web.json_response({"error": "비활성 계정"}, status=401)
    request["user"] = user
    return await handler(request)


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  설정 관리
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

class Config:
    def __init__(self):
        self.data = dict(DEFAULT_CONFIG)
        self.load()
    def load(self):
        try:
            if os.path.exists(CONFIG_FILE):
                with open(CONFIG_FILE, "r") as f:
                    self.data.update(json.load(f))
        except Exception:
            pass
    def save(self):
        with open(CONFIG_FILE, "w") as f:
            json.dump(self.data, f, indent=2, ensure_ascii=False)
    def get(self, key):
        return self.data.get(key, DEFAULT_CONFIG.get(key))
    def update(self, updates):
        self.data.update(updates)
        self.save()

config = Config()


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  연결 추적
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

@dataclass
class Connection:
    charger_id: str
    connected_at: str
    target_url: str
    messages_in: int = 0
    messages_out: int = 0
    bytes_in: int = 0
    bytes_out: int = 0
    last_activity: str = ""
    status: str = "connected"
    last_action: str = ""

active_connections: Dict[str, Connection] = {}
connection_history: list = []
stats = {"total_connections": 0, "total_messages": 0, "total_bytes": 0,
         "start_time": datetime.now(timezone.utc).isoformat()}


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  OCPP 파싱 + 로깅
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def parse_ocpp(raw):
    try:
        msg = json.loads(raw)
        if not isinstance(msg, list) or len(msg) < 3:
            return {"type": "unknown", "raw": raw[:200]}
        result = {"type": MSG_TYPES.get(msg[0], f"TYPE_{msg[0]}"), "id": msg[1]}
        if msg[0] == 2:
            result["action"] = msg[2]
            result["payload"] = msg[3] if len(msg) > 3 else {}
        elif msg[0] == 3:
            result["payload"] = msg[2] if len(msg) > 2 else {}
        elif msg[0] == 4:
            result["error_code"] = msg[2] if len(msg) > 2 else ""
            result["error_desc"] = msg[3] if len(msg) > 3 else ""
        return result
    except json.JSONDecodeError:
        return {"type": "parse_error", "raw": raw[:200]}

os.makedirs(LOG_DIR, exist_ok=True)

def _log_message(charger_id, direction, raw, parsed):
    now = datetime.now(timezone.utc)
    ts = now.strftime("%Y-%m-%d %H:%M:%S.%f")[:-3]
    action = parsed.get("action", parsed.get("type", "?"))
    msg_type = parsed.get("type", "?")

    arrow = f"{G}▶ CP→CS{E}" if direction == "charger_to_server" else f"{C}◀ CS→CP{E}"
    print(f"{DIM}{ts}{E} {arrow} {Y}{charger_id}{E} [{msg_type}] {M}{action}{E}")

    if not config.get("log_enabled"):
        return

    date_str = now.strftime("%Y%m%d")
    base = f"{LOG_DIR}/{charger_id}_{date_str}"

    if config.get("log_format") in ("jsonl", "both"):
        entry = {"timestamp": now.isoformat(), "charger_id": charger_id, "direction": direction,
                 "message_type": msg_type, "action": parsed.get("action"),
                 "message_id": parsed.get("id"), "payload": parsed.get("payload", {}), "raw": raw}
        with open(f"{base}.jsonl", "a", encoding="utf-8") as f:
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")

    if config.get("log_format") in ("txt", "both"):
        dir_label = "CP→CS" if direction == "charger_to_server" else "CS→CP"
        with open(f"{base}.txt", "a", encoding="utf-8") as f:
            f.write(f"\n{'='*70}\n[{ts}] {dir_label} | {msg_type} | {action}\n{'─'*70}\n")
            try:
                f.write(json.dumps(json.loads(raw), indent=2, ensure_ascii=False) + "\n")
            except Exception:
                f.write(raw + "\n")

    if config.get("pcap_enabled"):
        _write_pcap(f"{base}.pcap", direction, raw, now)

def _write_pcap(filepath, direction, data, ts):
    raw_bytes = data.encode("utf-8")
    if not os.path.exists(filepath):
        with open(filepath, "wb") as f:
            f.write(struct.pack("<IHHIIII", 0xa1b2c3d4, 2, 4, 0, 0, 65535, 147))
    direction_byte = b'\x00' if direction == "charger_to_server" else b'\x01'
    payload = direction_byte + raw_bytes
    sec = int(ts.timestamp())
    usec = int((ts.timestamp() - sec) * 1_000_000)
    with open(filepath, "ab") as f:
        f.write(struct.pack("<IIII", sec, usec, len(payload), len(payload)))
        f.write(payload)

# 실시간 브로드캐스트
live_ws_clients: Set = set()

def log_message(charger_id, direction, raw, parsed):
    _log_message(charger_id, direction, raw, parsed)
    entry = {"timestamp": datetime.now(timezone.utc).isoformat(), "charger_id": charger_id,
             "direction": direction, "type": parsed.get("type", "?"),
             "action": parsed.get("action", ""),
             "payload_preview": json.dumps(parsed.get("payload", {}), ensure_ascii=False)[:300]}
    asyncio.ensure_future(broadcast_log(entry))

async def broadcast_log(entry):
    if not live_ws_clients:
        return
    data = json.dumps(entry, ensure_ascii=False, default=str)
    closed = set()
    for ws in live_ws_clients:
        try:
            await ws.send_str(data)
        except Exception:
            closed.add(ws)
    live_ws_clients -= closed


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  WebSocket 프록시
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

async def proxy_handler(ws_charger, path):
    path_parts = path.strip("/").split("/")
    charger_id = path_parts[-1] if path_parts else "unknown"
    target_base = config.get("target_url").rstrip("/")
    target_url = f"{target_base}{path}"
    subprotocol = config.get("ocpp_subprotocol")

    now_str = datetime.now(timezone.utc).isoformat()
    conn = Connection(charger_id=charger_id, connected_at=now_str,
                      target_url=target_url, last_activity=now_str)
    active_connections[charger_id] = conn
    stats["total_connections"] += 1

    logger.info(f"{G}충전기 연결: {charger_id} → {target_url}{E}")

    try:
        async with websockets.connect(
            target_url, subprotocols=[subprotocol] if subprotocol else None,
            ping_interval=30, ping_timeout=20, close_timeout=10, max_size=2**20,
        ) as ws_server:
            conn.status = "active"

            async def c2s():
                try:
                    async for msg in ws_charger:
                        parsed = parse_ocpp(msg)
                        log_message(charger_id, "charger_to_server", msg, parsed)
                        conn.messages_in += 1; conn.bytes_in += len(msg)
                        conn.last_activity = datetime.now(timezone.utc).isoformat()
                        conn.last_action = parsed.get("action", "")
                        stats["total_messages"] += 1; stats["total_bytes"] += len(msg)
                        await ws_server.send(msg)
                except websockets.exceptions.ConnectionClosed:
                    pass

            async def s2c():
                try:
                    async for msg in ws_server:
                        parsed = parse_ocpp(msg)
                        log_message(charger_id, "server_to_charger", msg, parsed)
                        conn.messages_out += 1; conn.bytes_out += len(msg)
                        conn.last_activity = datetime.now(timezone.utc).isoformat()
                        stats["total_messages"] += 1; stats["total_bytes"] += len(msg)
                        await ws_charger.send(msg)
                except websockets.exceptions.ConnectionClosed:
                    pass

            done, pending = await asyncio.wait(
                [asyncio.create_task(c2s()), asyncio.create_task(s2c())],
                return_when=asyncio.FIRST_COMPLETED)
            for t in pending:
                t.cancel()
    except Exception as e:
        logger.error(f"{R}프록시 에러 [{charger_id}]: {e}{E}")
        conn.status = "error"
    finally:
        conn.status = "disconnected"
        connection_history.append(asdict(conn))
        if len(connection_history) > 500:
            connection_history.pop(0)
        active_connections.pop(charger_id, None)
        logger.info(f"{Y}연결 해제: {charger_id} (in:{conn.messages_in} out:{conn.messages_out}){E}")


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  관리 웹 API
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

async def h_index(request):
    html_path = Path(__file__).parent / "proxy_control.html"
    if html_path.exists():
        return web.FileResponse(html_path)
    return web.Response(text="proxy_control.html not found", status=404)

# ── 인증 ──

async def h_login(request):
    data = await request.json()
    user = user_mgr.verify(data.get("username", ""), data.get("password", ""))
    if not user:
        return web.json_response({"error": "아이디 또는 비밀번호가 틀립니다"}, status=401)
    token = _create_token(user["username"], user["username"], user["role"])
    return web.json_response({
        "token": token,
        "user": {"username": user["username"], "display_name": user["display_name"], "role": user["role"]},
    })

async def h_me(request):
    u = request["user"]
    return web.json_response({"username": u["username"], "display_name": u["display_name"], "role": u["role"]})

# ── 사용자 관리 ──

async def h_users_list(request):
    if request["user"]["role"] != "admin":
        return web.json_response({"error": "관리자 권한 필요"}, status=403)
    return web.json_response({"users": user_mgr.list_users()})

async def h_users_create(request):
    if request["user"]["role"] != "admin":
        return web.json_response({"error": "관리자 권한 필요"}, status=403)
    data = await request.json()
    username = data.get("username", "").strip()
    password = data.get("password", "")
    if not username or not password:
        return web.json_response({"error": "아이디와 비밀번호 필수"}, status=400)
    if username in user_mgr.users:
        return web.json_response({"error": "이미 존재하는 아이디"}, status=409)
    user_mgr.add_user(username, password, data.get("display_name", ""), data.get("role", "viewer"))
    logger.info(f"사용자 생성: {username} by {request['user']['username']}")
    return web.json_response({"status": "ok", "username": username})

async def h_users_update(request):
    if request["user"]["role"] != "admin":
        return web.json_response({"error": "관리자 권한 필요"}, status=403)
    username = request.match_info["username"]
    data = await request.json()
    if user_mgr.update_user(username, data):
        return web.json_response({"status": "ok"})
    return web.json_response({"error": "사용자 없음"}, status=404)

async def h_users_delete(request):
    if request["user"]["role"] != "admin":
        return web.json_response({"error": "관리자 권한 필요"}, status=403)
    username = request.match_info["username"]
    if username == request["user"]["username"]:
        return web.json_response({"error": "자기 자신 삭제 불가"}, status=400)
    if user_mgr.delete_user(username):
        return web.json_response({"status": "deleted"})
    return web.json_response({"error": "사용자 없음"}, status=404)

# ── 프록시 상태/설정 ──

async def h_status(request):
    return web.json_response({
        "proxy_port": PROXY_PORT, "target_url": config.get("target_url"),
        "target_name": config.get("target_name"),
        "log_enabled": config.get("log_enabled"), "log_format": config.get("log_format"),
        "pcap_enabled": config.get("pcap_enabled"),
        "active_connections": {k: asdict(v) for k, v in active_connections.items()},
        "active_count": len(active_connections), "stats": stats,
    })

async def h_config_get(request):
    return web.json_response(config.data)

async def h_config_set(request):
    data = await request.json()
    config.update(data)
    logger.info(f"설정 변경 by {request['user']['username']}: {list(data.keys())}")
    return web.json_response({"status": "ok", "config": config.data})

async def h_connections(request):
    return web.json_response({
        "active": {k: asdict(v) for k, v in active_connections.items()},
        "history": connection_history[-50:],
    })

# ── 로그 ──

async def h_logs_list(request):
    files = []
    if os.path.exists(LOG_DIR):
        for f in sorted(os.listdir(LOG_DIR), reverse=True):
            fp = os.path.join(LOG_DIR, f)
            files.append({"name": f, "size": os.path.getsize(fp),
                          "modified": datetime.fromtimestamp(os.path.getmtime(fp)).isoformat()})
    return web.json_response({"files": files})

async def h_log_content(request):
    name = request.match_info["name"]
    filepath = os.path.join(LOG_DIR, name)
    if not os.path.exists(filepath):
        return web.json_response({"error": "파일 없음"}, status=404)
    if name.endswith(".jsonl"):
        lines = []
        with open(filepath, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    try: lines.append(json.loads(line))
                    except: pass
        return web.json_response({"entries": lines[-200:], "total": len(lines)})
    if name.endswith(".txt"):
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read()
        if len(content) > 50000:
            content = "...(앞부분 생략)...\n\n" + content[-50000:]
        return web.json_response({"content": content})
    return web.FileResponse(filepath)

async def h_log_download(request):
    name = request.match_info["name"]
    filepath = os.path.join(LOG_DIR, name)
    if not os.path.exists(filepath):
        return web.Response(text="없음", status=404)
    return web.FileResponse(filepath, headers={"Content-Disposition": f'attachment; filename="{name}"'})

async def h_log_delete(request):
    name = request.match_info["name"]
    filepath = os.path.join(LOG_DIR, name)
    if os.path.exists(filepath):
        os.remove(filepath)
        return web.json_response({"status": "deleted"})
    return web.json_response({"error": "없음"}, status=404)

# ── 실시간 로그 WebSocket (인증은 쿼리 파라미터) ──

async def h_ws_live(request):
    token = request.query.get("token")
    if token:
        payload = _decode_token(token)
        if not payload:
            return web.Response(text="인증 실패", status=401)
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    live_ws_clients.add(ws)
    try:
        async for msg in ws:
            pass
    finally:
        live_ws_clients.discard(ws)
    return ws


def create_web_app():
    app = web.Application(middlewares=[auth_middleware])
    app.router.add_get("/", h_index)
    # 인증
    app.router.add_post("/api/auth/login", h_login)
    app.router.add_get("/api/auth/me", h_me)
    # 사용자 관리
    app.router.add_get("/api/users", h_users_list)
    app.router.add_post("/api/users", h_users_create)
    app.router.add_put("/api/users/{username}", h_users_update)
    app.router.add_delete("/api/users/{username}", h_users_delete)
    # 상태/설정
    app.router.add_get("/api/status", h_status)
    app.router.add_get("/api/config", h_config_get)
    app.router.add_post("/api/config", h_config_set)
    app.router.add_get("/api/connections", h_connections)
    # 로그
    app.router.add_get("/api/logs", h_logs_list)
    app.router.add_get("/api/logs/{name}", h_log_content)
    app.router.add_get("/api/logs/{name}/download", h_log_download)
    app.router.add_delete("/api/logs/{name}", h_log_delete)
    # 실시간
    app.router.add_get("/ws/live", h_ws_live)
    return app


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
#  메인
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

async def main():
    target = config.get("target_url")
    print(f"""
{C}╔═══════════════════════════════════════════════════╗
║     OCPP 양방향 프록시 / 로깅 서버 v2.0           ║
║     (인증 + 사용자 관리 포함)                      ║
╠═══════════════════════════════════════════════════╣
║  프록시 포트  : {PROXY_PORT}                                ║
║  관리 웹     : {WEB_PORT} (로그인 필요)                    ║
║  타겟 서버   : {target:<37s}║
║  초기 계정   : admin / admin1234                    ║
╚═══════════════════════════════════════════════════╝{E}
""")
    ws_server = await websockets.serve(
        proxy_handler, "0.0.0.0", PROXY_PORT,
        subprotocols=["ocpp1.6", "ocpp2.0.1"],
        ping_interval=30, ping_timeout=20, max_size=2**20)
    logger.info(f"프록시 서버: 포트 {PROXY_PORT}")

    web_app = create_web_app()
    runner = web.AppRunner(web_app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", WEB_PORT)
    await site.start()
    logger.info(f"관리 웹서버: 포트 {WEB_PORT}")
    logger.info(f"{G}서버 대기 중... (Ctrl+C 종료){E}")
    await asyncio.Future()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("서버 종료")