diff --git a/app/scripts/script1.py b/app/scripts/script1.py index 2d02dd8..903f2c7 100644 --- a/app/scripts/script1.py +++ b/app/scripts/script1.py @@ -8,6 +8,9 @@ import logging import json import urllib3 import re +import os +import hvac + urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning) logging.basicConfig( level=logging.INFO, @@ -17,12 +20,43 @@ logging.basicConfig( class LoginError(Exception): """Raised when login fails due to invalid credentials or missing CSRF token.""" + +def setup_vault(): + """Sets up and returns a Vault client.""" + vault_addr = os.environ.get("VAULT_ADDR", "https://vault.sectorq.eu") + try: + value_token = os.environ.get("VAULT_TOKEN") + if value_token is None: + raise KeyError + except KeyError: + value_token = prompt("Valult root token : ", is_password=True) + print("Use command export VAULT_TOKEN=") + os.environ["VAULT_TOKEN"] = value_token + client = hvac.Client(url=vault_addr, token=value_token) + + # Check if connected + if client.is_authenticated(): + print("Connected to Vault") + else: + print("Failed to login") + sys.exit(1) + return client + + +def get_secret(client, path, field=None): + """Retrieves a secret from Vault at the specified path.""" + return client.secrets.kv.v2.read_secret_version( + path=path, mount_point="secret", raise_on_deleted_version=True + )["data"]["data"][field] + def get_token(args): """Get CSRF token and cookies""" + vault_client = setup_vault() + unifi_password = get_secret(vault_client, "unifi/password", field="value") print("Logging in to Unifi Controller...") session = requests.Session() login_url = "https://unifi.sectorq.eu/api/auth/login" - payload_login = {"username": "jaydee", "password": "l4c1j4yd33Du5lo"} + payload_login = {"username": "jaydee", "password": unifi_password} response = session.post(login_url, json=payload_login, verify=False) response.raise_for_status() @@ -35,6 +69,7 @@ def get_token(args): def ban_ip(args): """Ban or unban IP address""" + data = get_token(args) logging.info(f"received payload: {args}") headers = {"x-csrf-token": data[0]}