#!/usr/bin/python

from ansible.module_utils.basic import AnsibleModule
import requests
import json
import os
import pickle

def get_portforwarding_rules(controller_url, username, password, site_id, mfa_token, cookies_file):
    # Disable SSL warnings
    requests.packages.urllib3.disable_warnings(requests.packages.urllib3.exceptions.InsecureRequestWarning)

    # Create session
    session = requests.Session()

    # Load cookies if available
    if os.path.exists(cookies_file):
        with open(cookies_file, 'rb') as f:
            session.cookies.update(pickle.load(f))

        # Test if the session is still valid by checking an endpoint
        test_url = f"{controller_url}/api/s/{site_id}/stat/sta"
        test_response = session.get(test_url, verify=False)
        if test_response.status_code != 200:
            os.remove(cookies_file)  # Remove expired cookies
    else:
        # Initial login URL
        login_url = f"{controller_url}/api/login"

        # Combine password and MFA token
        combined_password = f"{password}|{mfa_token}"

        # Debugging: Print combined password
        print(f"Combined Password: {combined_password}")

        # Login payload
        login_data = {
            "username": username,
            "password": combined_password
        }

        # Login headers
        headers = {
            "Content-Type": "application/json"
        }

        # Attempt to log in
        response = session.post(login_url, json=login_data, headers=headers, verify=False)

        if response.status_code == 200:
            # Save cookies to file
            with open(cookies_file, 'wb') as f:
                pickle.dump(session.cookies, f)
        else:
            return {"failed": True, "msg": "Login failed.", "response": response.json()}

    # API endpoint for fetching port forwarding rules
    portforwarding_url = f"{controller_url}/api/s/{site_id}/rest/portforward"

    # Fetch the port forwarding rules
    portforwarding_response = session.get(portforwarding_url, verify=False)
    portforwarding_response.raise_for_status()

    portforwarding_rules = portforwarding_response.json()
    return {"failed": False, "portforwarding_rules": portforwarding_rules}

def main():
    module_args = dict(
        controller_url=dict(type='str', required=True),
        username=dict(type='str', required=True),
        password=dict(type='str', required=True),
        site_id=dict(type='str', required=True),
        mfa_token=dict(type='str', required=True),
        cookies_file=dict(type='str', default="unifi_session_cookies.pkl")
    )

    module = AnsibleModule(argument_spec=module_args, supports_check_mode=True)

    result = get_portforwarding_rules(
        module.params['controller_url'],
        module.params['username'],
        module.params['password'],
        module.params['site_id'],
        module.params['mfa_token'],
        module.params['cookies_file']
    )

    if result['failed']:
        module.fail_json(msg=result['msg'], response=result.get('response', None))
    else:
        module.exit_json(changed=False, portforwarding_rules=result['portforwarding_rules'])

if __name__ == '__main__':
    main()