First Round Async. Much work left to do, dont trust results of hash categorization presently.
This commit is contained in:
+14
-27
@@ -5,6 +5,8 @@ import sys
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional, TypeVar
|
||||
|
||||
import aiofiles
|
||||
|
||||
T = TypeVar("T")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -19,21 +21,20 @@ PROTECTED_KEYS = [
|
||||
|
||||
_protected_config = {}
|
||||
|
||||
def get_system_config_path() -> Path:
|
||||
# Check inside bundled EXE directory first
|
||||
async def get_system_config_path() -> Path:
|
||||
bundled_dir = Path(getattr(sys, '_MEIPASS', ''))
|
||||
bundled_path = bundled_dir / "system_config.json"
|
||||
if bundled_path.exists():
|
||||
return bundled_path
|
||||
|
||||
# Fallback to external location
|
||||
return Path(__file__).parent.parent / "system_config.json"
|
||||
|
||||
def load_protected_config() -> dict:
|
||||
async def load_protected_config() -> dict:
|
||||
global _protected_config
|
||||
try:
|
||||
with open(get_system_config_path(), "r") as f:
|
||||
system_config = json.load(f)
|
||||
config_path = await get_system_config_path()
|
||||
async with aiofiles.open(config_path, "r") as f:
|
||||
content = await f.read()
|
||||
system_config = json.loads(content)
|
||||
except FileNotFoundError:
|
||||
logging.warning("⚠️ system_config.json not found. Using built-in defaults.")
|
||||
system_config = {
|
||||
@@ -46,10 +47,10 @@ def load_protected_config() -> dict:
|
||||
}
|
||||
}
|
||||
|
||||
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS}
|
||||
_protected_config = {key: system_config[key] for key in PROTECTED_KEYS if key in system_config}
|
||||
return _protected_config
|
||||
|
||||
def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
async def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
value = _protected_config.get(key)
|
||||
if value is None:
|
||||
logging.warning(f"Protected config key '{key}' not found.")
|
||||
@@ -62,7 +63,7 @@ def get_protected_value(key: str, cast_type: Callable[[str], T] = str, default:
|
||||
logging.warning(f"Invalid value for protected key '{key}': {value}. Expected type {cast_type.__name__}.")
|
||||
return default
|
||||
|
||||
def get_protected_json(key: str, default: str = "{}") -> dict:
|
||||
async def get_protected_json(key: str, default: str = "{}") -> dict:
|
||||
raw = _protected_config.get(key, default)
|
||||
if isinstance(raw, dict):
|
||||
return raw
|
||||
@@ -75,11 +76,8 @@ def get_protected_json(key: str, default: str = "{}") -> dict:
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to parse protected JSON key '{key}': {e}")
|
||||
return json.loads(default)
|
||||
|
||||
|
||||
|
||||
|
||||
def load_env_json(key: str, default: str):
|
||||
async def load_env_json(key: str, default: str):
|
||||
raw = os.getenv(key, default)
|
||||
try:
|
||||
return json.loads(raw)
|
||||
@@ -91,24 +89,13 @@ def load_env_json(key: str, default: str):
|
||||
logging.error(f"Failed to parse {key}: {e}")
|
||||
return json.loads(default)
|
||||
|
||||
def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
"""
|
||||
Safely retrieves an environment variable and casts it to the desired type.
|
||||
|
||||
Parameters:
|
||||
key (str): The name of the environment variable.
|
||||
cast_type (Callable[[str], T], optional): Function to cast the value. Defaults to str.
|
||||
default (Optional[T], optional): Default value if the variable is not set or invalid.
|
||||
|
||||
Returns:
|
||||
Optional[T]: The casted value or the default.
|
||||
"""
|
||||
async def load_env(key: str, cast_type: Callable[[str], T] = str, default: Optional[T] = None) -> Optional[T]:
|
||||
value = os.getenv(key)
|
||||
if value is None:
|
||||
logger.warning(f"Environment variable '{key}' not set.")
|
||||
return default
|
||||
try:
|
||||
value = value.strip("'\"") # Strip surrounding quotes
|
||||
value = value.strip("'\"")
|
||||
return cast_type(value)
|
||||
except (ValueError, TypeError):
|
||||
logger.warning(f"Invalid value for env var '{key}': {value}. Expected type {cast_type.__name__}.")
|
||||
|
||||
+50
-48
@@ -20,7 +20,7 @@ import re
|
||||
import dotenv
|
||||
import pandas as pd
|
||||
|
||||
import services.policyhandler as policyh
|
||||
import services.PolicyHandler as policyh
|
||||
from flows.otp import generate, otp_activities_by_agent, revoke
|
||||
from flows.prepPolicy import (
|
||||
buildPathsandPublishers,
|
||||
@@ -32,8 +32,9 @@ from flows.prepPolicy import (
|
||||
from flows.quietAgent import findQuietAgents
|
||||
from services.agenthandler import findAgents, moveAgentToRelatedPolicy, selectAgents
|
||||
from services.API import AirlockAPIWrapper
|
||||
from services.TaskQueue import AsyncTaskQueue
|
||||
from utils.configmanager import load_env
|
||||
from utils.selector import Selector
|
||||
from utils.Selector import Selector
|
||||
from utils.utils import (
|
||||
areYouSure,
|
||||
colorText,
|
||||
@@ -47,9 +48,9 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
dotenv.load_dotenv()
|
||||
|
||||
def menu_main(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
extras = load_env("EXTRAS")
|
||||
async def menu_main(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
|
||||
working_dir = await load_env("WORKING_DIR")
|
||||
extras = await load_env("EXTRAS")
|
||||
while True:
|
||||
displayIntro()
|
||||
# Add Settings, and give option to change working dir
|
||||
@@ -63,34 +64,34 @@ def menu_main(api: AirlockAPIWrapper):
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("Q. 🔚 - Quit", "yellow"))
|
||||
|
||||
choice = get_sanitized_input("\nEnter Menu Item: ")
|
||||
choice = await get_sanitized_input("\nEnter Menu Item: ")
|
||||
if choice == "1":
|
||||
print("This Feature is still in development")
|
||||
get_sanitized_input("Press enter to continue")
|
||||
await get_sanitized_input("Press enter to continue")
|
||||
elif choice == "2":
|
||||
menu_otp(api)
|
||||
await menu_otp(api)
|
||||
elif choice == "3":
|
||||
choices = ["audit", "enforcement"]
|
||||
print(colorText("Move devices to which state?:", "yellow"))
|
||||
direction = Selector.select_string(choices, False, False)
|
||||
devices = selectAgents(api)
|
||||
direction = await Selector.select_string(choices, False, False)
|
||||
devices = await selectAgents(api)
|
||||
print(colorText("Would you like to continue with these devices?","white"))
|
||||
for device in devices:
|
||||
print(device.hostname)
|
||||
confirm = Selector.confirm()
|
||||
confirm =await Selector.confirm()
|
||||
if direction and devices and confirm:
|
||||
for device in devices:
|
||||
moveAgentToRelatedPolicy(api,device, direction[0])
|
||||
await moveAgentToRelatedPolicy(api,device, direction)
|
||||
elif choice == "4":
|
||||
findAgents(api,False)
|
||||
await findAgents(api,False)
|
||||
elif choice == "5":
|
||||
findQuietAgents(api)
|
||||
await findQuietAgents(api)
|
||||
elif choice == "6":
|
||||
if extras == "POLICYPREP": menu_policymanagment(api)
|
||||
if extras == "POLICYPREP": await menu_policymanagment(api, queue)
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
await open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
await menu_settings()
|
||||
elif choice.upper() == "Q":
|
||||
break
|
||||
else:
|
||||
@@ -98,7 +99,7 @@ def menu_main(api: AirlockAPIWrapper):
|
||||
|
||||
|
||||
|
||||
def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
async def menu_policy_enforce(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
|
||||
selected_policies = []
|
||||
destination_policy = []
|
||||
destination_allowlist = []
|
||||
@@ -106,34 +107,35 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
processed_hashes = []
|
||||
processed_publishers = []
|
||||
tested = False
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
working_dir = await load_env("WORKING_DIR")
|
||||
|
||||
while True:
|
||||
printEnforceChecklist(selected_policies, destination_policy, destination_allowlist)
|
||||
choice = get_sanitized_input("\nEnter your choice: ")
|
||||
choice = await get_sanitized_input("\nEnter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
selected_policies = selectPolicies(api,True)
|
||||
selected_policies = await selectPolicies(api,True)
|
||||
|
||||
elif choice == "2":
|
||||
print(colorText("Please choose destination_name Policy for Path Exclusions", "white"))
|
||||
|
||||
destination_policy = selectPolicies(api, False)
|
||||
destination_policy = await selectPolicies(api, False)
|
||||
|
||||
print(colorText("Please choose Allowlist for Hashes", "white"))
|
||||
|
||||
destination_allowlist = selectAllowlists(api, destination_policy, False)
|
||||
destination_allowlist = await selectAllowlists(api, destination_policy, False) # pyright: ignore[reportArgumentType]
|
||||
|
||||
elif choice == "3":
|
||||
sortHashes(
|
||||
await sortHashes(
|
||||
api,
|
||||
queue,
|
||||
selected_policies,
|
||||
type=[1, 2, 6, 7],
|
||||
)
|
||||
|
||||
elif choice == "4":
|
||||
if os.path.exists(f"{working_dir}\\Needs_Review\\Review_First\\approved_executions.csv"):
|
||||
buildPathsandPublishers(False)
|
||||
await buildPathsandPublishers(False)
|
||||
else:
|
||||
print("File not found. Please make sure it's saved correctly and try again.")
|
||||
|
||||
@@ -141,7 +143,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
if os.path.exists(f"{working_dir}\\Approved\\hashes_to_add.csv") and os.path.exists(
|
||||
f"{working_dir}\\Approved\\primary_Paths.csv"
|
||||
):
|
||||
buildPreflights()
|
||||
await buildPreflights()
|
||||
else:
|
||||
print("File not found. Please make sure it's saved correctly and try again.")
|
||||
|
||||
@@ -209,7 +211,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
|
||||
elif choice == "7":
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
confirmation = await get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if (
|
||||
tested
|
||||
and destination_policy
|
||||
@@ -217,10 +219,10 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
and confirmation.strip() == "I AGREE"
|
||||
):
|
||||
print(colorText("Proceeding with the code...", "yellow"))
|
||||
api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
|
||||
api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
|
||||
await api.hash_add_to_allowlist(destination_allowlist[0].applicationid, processed_hashes)
|
||||
await api.policy_add_path_exclusions(destination_policy[0].groupid, processed_paths)
|
||||
if processed_publishers:
|
||||
api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
|
||||
await api.policy_add_publishers(destination_policy[0].groupid, processed_publishers)
|
||||
else:
|
||||
logger.error("Confirmation block failed. Reasons:")
|
||||
if not tested:
|
||||
@@ -233,9 +235,9 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
logger.error(" - User did not confirm with 'I AGREE'. Received: '%s'", confirmation.strip())
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
await open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
await menu_settings()
|
||||
elif choice.upper == "B":
|
||||
break
|
||||
|
||||
@@ -244,7 +246,7 @@ def menu_policy_enforce(api: AirlockAPIWrapper):
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
|
||||
|
||||
def menu_otp(api: AirlockAPIWrapper):
|
||||
async def menu_otp(api: AirlockAPIWrapper):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
while True:
|
||||
|
||||
@@ -256,23 +258,23 @@ def menu_otp(api: AirlockAPIWrapper):
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
|
||||
choice = get_sanitized_input("Enter your choice: ")
|
||||
choice = await get_sanitized_input("Enter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
otp_list = generate(api)
|
||||
otp_list = await generate(api)
|
||||
print(colorText(otp_list,"green"))
|
||||
elif choice == "2":
|
||||
otp_activities_by_agent(api)
|
||||
await otp_activities_by_agent(api)
|
||||
elif choice == "3":
|
||||
revoke(api)
|
||||
await revoke(api)
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
await open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
await menu_settings()
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
|
||||
def menu_policymanagment(api: AirlockAPIWrapper):
|
||||
async def menu_policymanagment(api: AirlockAPIWrapper, queue: AsyncTaskQueue):
|
||||
working_dir = load_env("WORKING_DIR")
|
||||
while True:
|
||||
print(colorText("1. 🔒 - Prepare Policy For Enforcement", "yellow"))
|
||||
@@ -280,31 +282,31 @@ def menu_policymanagment(api: AirlockAPIWrapper):
|
||||
print(colorText("F. 📂 - Open Working Directory", "yellow"))
|
||||
print(colorText("S. 🛠️ - Settings", "yellow"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
choice = get_sanitized_input("\n Enter Menu Item: ")
|
||||
choice = await get_sanitized_input("\n Enter Menu Item: ")
|
||||
|
||||
if choice == "1":
|
||||
menu_policy_enforce(api)
|
||||
await menu_policy_enforce(api, queue)
|
||||
elif choice == "2":
|
||||
areYouSure()
|
||||
confirmation = get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
confirmation = await get_sanitized_input("Type 'I AGREE' to continue: ")
|
||||
if confirmation.strip() == "I AGREE":
|
||||
policyh.updateAuditPoliciesFromEnforcementPolices(api)
|
||||
await policyh.updateAuditPoliciesFromEnforcementPolices(api)
|
||||
|
||||
elif choice.upper() == "F":
|
||||
open_directory(working_dir)
|
||||
await open_directory(working_dir)
|
||||
elif choice.upper() == "S":
|
||||
menu_settings()
|
||||
await menu_settings()
|
||||
elif choice.upper() == "B":
|
||||
break
|
||||
else:
|
||||
print(colorText("Invalid choice. Please try again.", "red"))
|
||||
def menu_settings():
|
||||
async def menu_settings():
|
||||
while True:
|
||||
print(colorText("\n--- 🛠️ Settings Submenu 🛠️ ---", "cyan"))
|
||||
print(colorText("This Feature is still in development", "cyan"))
|
||||
# print(colorText("2. Sub-option B","cyan"))
|
||||
print(colorText("B. 🔙 - Back", "yellow"))
|
||||
choice = get_sanitized_input("Enter your choice: ")
|
||||
choice = await get_sanitized_input("Enter your choice: ")
|
||||
|
||||
if choice == "1":
|
||||
pass #TODO ADD CHANGE WORKDIR CODE
|
||||
|
||||
+20
-54
@@ -1,33 +1,16 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
import logging
|
||||
from typing import Any, Callable, List, Optional, Union
|
||||
|
||||
from utils.utils import get_sanitized_input
|
||||
from utils.utils import get_sanitized_input # new async version
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Selector:
|
||||
@staticmethod
|
||||
def _display_choices(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
num_columns: int = 4,
|
||||
num_columns: int = 3,
|
||||
header: str = "Available Choices:"
|
||||
) -> None:
|
||||
sorted_items = sorted(items, key=lambda item: label_func(item).lower())
|
||||
@@ -43,7 +26,7 @@ class Selector:
|
||||
print(line)
|
||||
|
||||
@staticmethod
|
||||
def _select_from_list(
|
||||
async def _select_from_list(
|
||||
items: List[Any],
|
||||
label_func: Callable[[Any], str],
|
||||
allow_multiple: bool = False,
|
||||
@@ -60,7 +43,7 @@ class Selector:
|
||||
|
||||
if allow_multiple:
|
||||
while True:
|
||||
choice = get_sanitized_input("Select an item by number (or Q to finish): ").strip().lower()
|
||||
choice = (await get_sanitized_input("Select an item by number (or Q to finish): ")).strip().lower()
|
||||
if choice == "q":
|
||||
break
|
||||
try:
|
||||
@@ -80,7 +63,7 @@ class Selector:
|
||||
return selected if selected else None
|
||||
else:
|
||||
try:
|
||||
choice = int(get_sanitized_input("Select one item by number: "))
|
||||
choice = int(await get_sanitized_input("Select one item by number: "))
|
||||
if 1 <= choice <= len(sorted_items):
|
||||
selected_item = sorted_items[choice - 1]
|
||||
logger.info(f"Selected: {label_func(selected_item)}")
|
||||
@@ -92,12 +75,8 @@ class Selector:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def select_objects(
|
||||
objects: List[Any],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[Any], List[Any]]:
|
||||
return Selector._select_from_list(
|
||||
async def select_objects(objects: List[Any], allow_multiple: bool = False, prompt_each: bool = False) -> Union[Optional[Any], List[Any]]:
|
||||
return await Selector._select_from_list(
|
||||
objects,
|
||||
label_func=lambda obj: getattr(obj, "name", str(obj)),
|
||||
allow_multiple=allow_multiple,
|
||||
@@ -106,12 +85,8 @@ class Selector:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_string(
|
||||
options: List[str],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[str], List[str]]:
|
||||
return Selector._select_from_list(
|
||||
async def select_string(options: List[str], allow_multiple: bool = False, prompt_each: bool = False) -> Union[Optional[str], List[str]]:
|
||||
return await Selector._select_from_list(
|
||||
options,
|
||||
label_func=str,
|
||||
allow_multiple=allow_multiple,
|
||||
@@ -120,12 +95,8 @@ class Selector:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_int(
|
||||
options: List[int],
|
||||
allow_multiple: bool = False,
|
||||
prompt_each: bool = False
|
||||
) -> Union[Optional[int], List[int]]:
|
||||
return Selector._select_from_list(
|
||||
async def select_int(options: List[int], allow_multiple: bool = False, prompt_each: bool = False) -> Union[Optional[int], List[int]]:
|
||||
return await Selector._select_from_list(
|
||||
options,
|
||||
label_func=lambda x: str(x),
|
||||
allow_multiple=allow_multiple,
|
||||
@@ -134,16 +105,11 @@ class Selector:
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def select_value(
|
||||
prompt: str,
|
||||
value_type: type = int,
|
||||
valid_range: Optional[tuple] = None,
|
||||
allow_quit: bool = False
|
||||
) -> Optional[Any]:
|
||||
async def select_value(prompt: str, value_type: type = int, valid_range: Optional[tuple] = None, allow_quit: bool = False) -> Optional[Any]:
|
||||
while True:
|
||||
user_input = get_sanitized_input(prompt).strip().lower()
|
||||
user_input = (await get_sanitized_input(prompt)).strip().lower()
|
||||
if allow_quit and user_input == "q":
|
||||
logger.info("User opted to quit value selection.")
|
||||
logger.debug("User opted to quit value selection.")
|
||||
return None
|
||||
try:
|
||||
value = value_type(user_input)
|
||||
@@ -152,20 +118,20 @@ class Selector:
|
||||
if not (min_val <= value <= max_val):
|
||||
logger.warning(f"Value out of range ({min_val}–{max_val}).")
|
||||
continue
|
||||
logger.info(f"User selected value: {value}")
|
||||
logger.debug(f"User selected value: {value}")
|
||||
return value
|
||||
except ValueError:
|
||||
logger.warning(f"Invalid input. Expected a {value_type.__name__}.")
|
||||
|
||||
@staticmethod
|
||||
def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
|
||||
async def confirm(prompt: str = "Are you sure? (Y/N): ") -> bool:
|
||||
while True:
|
||||
response = get_sanitized_input(prompt).strip().lower()
|
||||
response = (await get_sanitized_input(prompt)).strip().lower()
|
||||
if response in ["y", "yes"]:
|
||||
logger.info("User confirmed action.")
|
||||
logger.debug("User confirmed action.")
|
||||
return True
|
||||
elif response in ["n", "no"]:
|
||||
logger.info("User declined action.")
|
||||
logger.debug("User declined action.")
|
||||
return False
|
||||
else:
|
||||
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
|
||||
logger.warning("Invalid confirmation input. Expected 'Y' or 'N'.")
|
||||
|
||||
+10
-24
@@ -1,17 +1,3 @@
|
||||
# Copyright (C) 2025 James Brotosky, Brandon Wickline
|
||||
#
|
||||
# This program is free software: you can redistribute it and/or modify
|
||||
# it under the terms of the GNU Affero General Public License as published
|
||||
# by the Free Software Foundation, either version 3 of the License, or
|
||||
# (at your option) any later version.
|
||||
#
|
||||
# This program is distributed in the hope that it will be useful,
|
||||
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
# GNU Affero General Public License for more details.
|
||||
#
|
||||
# You should have received a copy of the GNU Affero General Public License
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
import json
|
||||
import logging
|
||||
@@ -41,7 +27,11 @@ def configure_logging(log_dir: Path, log_level: str = "DEBUG"):
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG))
|
||||
|
||||
# 🔧 Clear existing handlers
|
||||
|
||||
httpx_logger = logging.getLogger("httpx")
|
||||
httpx_logger.setLevel(getattr(logging, log_level.upper(), logging.DEBUG))
|
||||
|
||||
|
||||
for handler in logger.handlers[:]:
|
||||
logger.removeHandler(handler)
|
||||
|
||||
@@ -104,14 +94,14 @@ def load_user_config(config_dir: Path) -> dict:
|
||||
def write_config_to_env(config: dict, env_path: Path):
|
||||
for key, value in config.items():
|
||||
if key in PROTECTED_KEYS:
|
||||
continue # Skip protected keys
|
||||
continue
|
||||
try:
|
||||
serialized = json.dumps(value) if isinstance(value, (list, dict)) else str(value)
|
||||
set_key(env_path, key, serialized)
|
||||
except Exception as e:
|
||||
logging.warning(f"Failed to write {key} to .env: {e}")
|
||||
|
||||
def setup() -> Path:
|
||||
async def setup() -> Path:
|
||||
base_dir = get_base_directory()
|
||||
dirs = {
|
||||
'config': base_dir / 'config',
|
||||
@@ -129,6 +119,7 @@ def setup() -> Path:
|
||||
env_path = base_dir / ".env"
|
||||
if not env_path.exists():
|
||||
env_path.touch()
|
||||
|
||||
load_dotenv(dotenv_path=env_path, override=True)
|
||||
|
||||
working_dir = Path(os.getenv("WORKING_DIR") or (base_dir / "data"))
|
||||
@@ -155,14 +146,10 @@ def setup() -> Path:
|
||||
|
||||
user_config = load_user_config(dirs['config'])
|
||||
merged_config = {**system_config, **user_config}
|
||||
|
||||
protected_config = load_protected_config()
|
||||
protected_config = await load_protected_config()
|
||||
merged_config.update(protected_config)
|
||||
|
||||
# ✅ URL resolution order: system_config → .env → user prompt
|
||||
url = system_config.get("URL")
|
||||
if not url:
|
||||
url = os.getenv("URL")
|
||||
url = system_config.get("URL") or os.getenv("URL")
|
||||
if not url:
|
||||
url = input("🌐 Enter the service URL (e.g., https://example.com/api): ").strip()
|
||||
merged_config["URL"] = url
|
||||
@@ -171,5 +158,4 @@ def setup() -> Path:
|
||||
logging.debug(f"Service URL set to: {url}")
|
||||
|
||||
write_config_to_env(merged_config, env_path)
|
||||
|
||||
return working_dir
|
||||
+10
-12
@@ -14,6 +14,7 @@
|
||||
# along with this program. If not, see <https://www.gnu.org/licenses/>.
|
||||
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
@@ -99,15 +100,15 @@ def choose_file(initial_directory=None, required_substring=None):
|
||||
|
||||
|
||||
|
||||
|
||||
def get_sanitized_input(prompt: str) -> str:
|
||||
async def get_sanitized_input(prompt: str) -> str:
|
||||
while True:
|
||||
user_input = input(prompt)
|
||||
user_input = await asyncio.to_thread(input, prompt)
|
||||
if user_input.strip() == "":
|
||||
return user_input # Allow blank lines
|
||||
if re.match(r'^[a-zA-Z0-9_\- .]+$', user_input.strip()):
|
||||
return user_input
|
||||
if re.match(r'^[a-zA-Z0-9_ .-]+$', user_input.strip()):
|
||||
return user_input
|
||||
else:
|
||||
logger.debug("User entered invalid input")
|
||||
print("Invalid input. Only letters, numbers, underscores, spaces, hyphens, and periods are allowed.")
|
||||
|
||||
|
||||
@@ -695,14 +696,11 @@ def formatHTML(df, output_html_path=None, overwrite=True):
|
||||
|
||||
|
||||
|
||||
def open_directory(path):
|
||||
async def open_directory(path):
|
||||
system = platform.system()
|
||||
|
||||
if system == "Windows":
|
||||
os.startfile(path)
|
||||
await asyncio.to_thread(os.startfile, path)
|
||||
elif system == "Linux":
|
||||
subprocess.run(["xdg-open", path])
|
||||
await asyncio.to_thread(subprocess.run, ["xdg-open", path])
|
||||
else:
|
||||
raise OSError(f"Unsupported operating system: {system}")
|
||||
|
||||
|
||||
raise OSError(f"Unsupported operating system: {system}")
|
||||
Reference in New Issue
Block a user