import requests
import yaml
import os
import sys
from requests.auth import HTTPBasicAuth
import platform

def load_config(config_path):
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)

def search_certificates(config):
    url = f"{config['host']}/api/trust/cert/search"
    payload = {
        "searchPhrase": config['certificate_search']
    }
    response = requests.post(url, json=payload, auth=HTTPBasicAuth(config['api_key'], config['api_secret']))
    response.raise_for_status()
    return response.json().get("rows", [])

def export_certificate(config, uuid):
    url = f"{config['host']}/api/trust/cert/generate_file/{uuid}/{config['export_format']}"
    response = requests.post(url, auth=HTTPBasicAuth(config['api_key'], config['api_secret']))
    response.raise_for_status()
    return response.json().get("payload")

def save_certificate(config, cert_data, filename, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    filepath = os.path.join(output_dir, filename)
    with open(filepath, 'w') as f:
        f.write(cert_data)
        
    if platform.system().lower() == "linux":
        try:
            import pwd
            import grp

            if 'file_mode' in config:
                os.chmod(filepath, int(config['file_mode'], 8))

            if 'file_owner' in config or 'file_group' in config:
                uid = pwd.getpwnam(config.get('file_owner', pwd.getpwuid(os.getuid()).pw_name)).pw_uid
                gid = grp.getgrnam(config.get('file_group', grp.getgrgid(os.getgid()).gr_name)).gr_gid
                os.chown(filepath, uid, gid)
        except Exception as e:
            print(f"[WARN] Error setting permissions: {e}")
    print(f"Certificate saved.")

def main():
    if len(sys.argv) != 2:
        print("Usage: python export_cert.py <config.yaml>")
        sys.exit(1)

    ### Load Config ###
    config = load_config(sys.argv[1])
    
    ### check cert options ###
    if config['export_format'] not in ['crt', 'prv']:
        print(f"Invalid export format. Possible options are crt or prv.")
        sys.exit(1)
    
    ### Search Certificates ###
    certificates = search_certificates(config)
    if len(certificates) > 1:
        print(f"Search results in more then one certificate. Please adjust your search to only return a single one.")
        sys.exit(1)
    if len(certificates) == 0:
        print(f"No certificate found with search phrase: {config['certificate_search']}")
        sys.exit(1)
        
    certificate_uuid = certificates[0]['uuid']
    
    ### Get Certificate by uuid ###
    certificate_data = export_certificate(config, certificate_uuid)
    save_certificate(config, certificate_data, 'cert.crt', config['output_directory'])

if __name__ == "__main__":
    main()