import socket
import struct
import hashlib
import random

# RADIUS Attribute Type names (RFC 2865 + common)
RADIUS_ATTR_NAMES = {
    1:  "User-Name",
    6:  "Service-Type",
    8:  "Framed-IP-Address",
    11: "Filter-Id",
    18: "Reply-Message",
    25: "Class",
    26: "Vendor-Specific",
    27: "Session-Timeout",
}

def encrypt_password(password: str, secret: str, request_authenticator: bytes) -> bytes:
    """Mã hóa mật khẩu PAP theo chuẩn RADIUS RFC 2865."""
    password_bytes = password.encode('utf-8')
    if len(password_bytes) % 16 != 0:
        password_bytes += b'\x00' * (16 - len(password_bytes) % 16)

    encrypted = b''
    last = secret.encode('utf-8') + request_authenticator
    for i in range(0, len(password_bytes), 16):
        block = password_bytes[i:i+16]
        md5_hash = hashlib.md5(last).digest()
        encrypted_block = bytes(a ^ b for a, b in zip(block, md5_hash))
        encrypted += encrypted_block
        last = secret.encode('utf-8') + encrypted_block
    return encrypted


def build_access_request(username: str, password: str, secret: str,
                         identifier: int = None) -> bytes:
    """Tạo gói Access-Request PAP cho RADIUS server."""
    if identifier is None:
        identifier = random.randint(0, 255)

    request_authenticator = bytes(random.getrandbits(8) for _ in range(16))

    username_attr = b'\x01' + struct.pack('B', len(username) + 2) + username.encode('utf-8')
    password_attr_bytes = encrypt_password(password, secret, request_authenticator)
    password_attr = b'\x02' + struct.pack('B', len(password_attr_bytes) + 2) + password_attr_bytes

    attrs = username_attr + password_attr
    length = 20 + len(attrs)
    header = struct.pack('!BBH', 1, identifier, length) + request_authenticator
    return header + attrs


def parse_radius_attributes(attributes_raw: bytes) -> dict:
    """
    Parse phần attributes của gói RADIUS response.
    Trả về dict: {type_int: [value_bytes, ...]}
    """
    result: dict[int, list] = {}
    idx = 0
    while idx < len(attributes_raw):
        if idx + 2 > len(attributes_raw):
            break
        attr_type = attributes_raw[idx]
        attr_len  = attributes_raw[idx + 1]
        if attr_len < 2 or idx + attr_len > len(attributes_raw):
            break
        attr_value = attributes_raw[idx + 2: idx + attr_len]
        result.setdefault(attr_type, []).append(attr_value)
        idx += attr_len
    return result


def check_radius_login(server: str, port: int, secret: str, username: str,
                       password: str, timeout: int = 5) -> bool:
    """
    Kiểm tra username/password với RADIUS server (PAP).
    Trả về True nếu Access-Accept, False nếu thất bại.
    (Giữ nguyên signature để không ảnh hưởng code cũ)
    """
    ok, _ = check_radius_login_extended(server, port, secret, username, password, timeout)
    return ok


def check_radius_login_extended(server: str, port: int, secret: str, username: str,
                                 password: str, timeout: int = 5) -> tuple[bool, dict]:
    """
    Kiểm tra username/password với RADIUS server (PAP).
    Trả về (success: bool, attributes: dict)
    attributes chứa các giá trị decode từ RADIUS response, ví dụ:
        {
            'filter_id': ['Trung tâm An ninh thông tin'],
            'reply_message': ['privacyIDEA access granted'],
            'raw': {11: [b'...'], 18: [b'...']},
        }
    """
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(timeout)
    attrs_decoded: dict = {}
    try:
        req_packet = build_access_request(username, password, secret)
        sock.sendto(req_packet, (server, port))

        resp, _ = sock.recvfrom(4096)
        code   = resp[0]
        length = struct.unpack('!H', resp[2:4])[0]
        raw_attrs = parse_radius_attributes(resp[20:length])

        # Decode các attributes thường gặp sang chuỗi
        def _decode_list(raw_list: list) -> list[str]:
            return [v.decode('utf-8', errors='replace') for v in raw_list]

        attrs_decoded['raw'] = raw_attrs

        # Filter-Id (Type 11) — thường chứa thông tin nhóm/đơn vị
        if 11 in raw_attrs:
            attrs_decoded['filter_id'] = _decode_list(raw_attrs[11])

        # Reply-Message (Type 18)
        if 18 in raw_attrs:
            attrs_decoded['reply_message'] = _decode_list(raw_attrs[18])

        # Class (Type 25)
        if 25 in raw_attrs:
            attrs_decoded['class'] = _decode_list(raw_attrs[25])

        if code == 2:  # Access-Accept
            print(f"[RADIUS] ✅ Đăng nhập thành công: {username}")
            if attrs_decoded.get('filter_id'):
                print(f"[RADIUS] Filter-Id: {attrs_decoded['filter_id']}")
            return True, attrs_decoded
        else:
            print(f"[RADIUS] ❌ Sai username hoặc password: {username}")
            return False, attrs_decoded

    except socket.timeout:
        print("[RADIUS] ⏱ Timeout khi kết nối đến RADIUS server.")
        return False, {}
    except Exception as e:
        print(f"[RADIUS] ⚠️ Lỗi: {e}")
        return False, {}
    finally:
        sock.close()