update

BackEnd/readtoml.py

@@ -1,73 +1,58 @@
 import subprocess
-from tomlkit import parse
-from tomlkit import dumps
-from tomlkit import table
+from tomlkit import loads
 import os
-# Read the configuration file
+
+# Set file paths
 test_directory = os.path.dirname(os.path.abspath(__file__))
 absolute_path = os.path.join(test_directory, '..', 'config', 'sampleconfig.toml')
-config_file_path = absolute_path 
-
+config_file_path = absolute_path
 rules_file_path = os.path.join(test_directory, '..', 'BackEnd', 'rules.conf')
 
 # Check if the configuration file exists
 if not os.path.exists(config_file_path):
     print("Error: Configuration file not found.")
     exit(1)
 
-
+# Define the parse_toml_file function
 def parse_toml_file(file_path):
     try:
         with open(file_path, 'r') as file:
             toml_content = file.read()
-            print(toml_content)
-            parsed_data = tomlkit.loads(toml_content)
+            parsed_data = loads(toml_content)
             return parsed_data
     except FileNotFoundError:
         print(f"Error: File not found at {file_path}")
         return None
-    except tomlkit.exceptions.ParseError as e:
-        print(f"Error parsing TOML file: {e}")
-        return None
 
-# Example usage
-toml_file_path = config_file_path  # Replace with the actual path to your TOML file
+# Parse the TOML configuration file
+parsed_data = parse_toml_file(config_file_path)
 
-parsed_data = parse_toml_file(toml_file_path)
-print(parsed_data)
+# Check if the 'enable' key is present and set to True
+enable = parsed_data.get('enable', False)
 
+# Check if USBGuard should be disabled
+if not enable:
+    subprocess.run(["sudo", "systemctl", "disable", "--now", "usbguard"])
+    exit()
 
+# Generate rules based on parsed data
+rules_content = ""
+allow_all = parsed_data.get('allow-all', False)
 
-''' 
-enable_usbguard = len(re.findall(r'enable\s*=\s*true', config_content, re.IGNORECASE))
-if enable_usbguard == 0:
-    subprocess.run(['sudo', 'systemctl', 'disable', '--now', 'readtoml'])
-    exit(0)
+if allow_all:
+    rules_content = "allow-all:\n  allow\n"
+else:
+    for rule in parsed_data.get('rules', []):
+        rules_content += f"  allow {rule['id']} name \"{rule['name']}\" via-port \"{rule['port']}\"\n"
 
-# Generate rules.conf
+# Write rules to rules.conf
 with open(rules_file_path, 'w') as rules_file:
-    rules_file.write("# USBGuard rules.conf\n")
-
-allow_all = len(re.findall(r'allow-all\s*=\s*true', config_content, re.IGNORECASE))
+    rules_file.write(rules_content)
 
-# Allow all or generate rules based on configuration
-if allow_all == 1:
-    with open('rules.conf', 'a') as rules_file:
-        rules_file.write("allow\n")
-else:
-    # Loop through each rule and add it to rules.conf
-    allow_rules = re.findall(r'^\s*allow\s*{.*?}', config_content, re.DOTALL | re.MULTILINE)
-    for rule in allow_rules:
-        id_value = re.search(r'id\s*=\s*"(.*?)"', rule).group(1)
-        name_value = re.search(r'name\s*=\s*"(.*?)"', rule).group(1)
-        port_value = re.search(r'port\s*=\s*"(.*?)"', rule).group(1)
-        with open('rules.conf', 'a') as rules_file:
-            rules_file.write(f'allow {id_value} name "{name_value}" via-port "{port_value}"\n')
+# Install rules
+subprocess.run(["sudo", "install", "-m", "0600", "-o", "root", "-g", "root", rules_file_path, "/etc/usbguard/rules.conf"])
 
-# Install rules and restart USBGuard
-subprocess.run(['sudo', 'install', '-m', '0600', '-o', 'root', '-g', 'root', 'rules.conf', '/etc/usbguard/rules.conf'])
-subprocess.run(['sudo', 'systemctl', 'restart', 'usbguard'])
-subprocess.run(['sudo', 'systemctl', 'enable', 'usbguard'])
+# Restart and enable USBGuard
+subprocess.run(["sudo", "systemctl", "restart", "usbguard"])
+subprocess.run(["sudo", "systemctl", "enable", "usbguard"])
 
-print("USBGuard configured successfully.")
-'''