more stuff
This commit is contained in:
617
data_loader.py
617
data_loader.py
@@ -1,276 +1,355 @@
|
||||
from models import Customer, CustomerService, Location, Host, Service, ServiceType, HostType, VPNType
|
||||
from typing import List
|
||||
import yaml
|
||||
from pathlib import Path
|
||||
from typing import List, Dict, Any
|
||||
from models import (
|
||||
Customer, CustomerService, Location, Host, Service,
|
||||
ServiceType, HostType, VPNType
|
||||
)
|
||||
|
||||
|
||||
def get_config_dir() -> Path:
|
||||
"""Get the VPNTray configuration directory path."""
|
||||
home = Path.home()
|
||||
config_dir = home / ".vpntray" / "customers"
|
||||
return config_dir
|
||||
|
||||
|
||||
def ensure_config_dir() -> Path:
|
||||
"""Ensure the configuration directory exists."""
|
||||
config_dir = get_config_dir()
|
||||
config_dir.mkdir(parents=True, exist_ok=True)
|
||||
return config_dir
|
||||
|
||||
|
||||
def parse_service_type(service_type_str: str) -> ServiceType:
|
||||
"""Convert a string to ServiceType enum, with fallback."""
|
||||
# Map common strings to enum values
|
||||
type_mapping = {
|
||||
"SSH": ServiceType.SSH,
|
||||
"Web GUI": ServiceType.WEB_GUI,
|
||||
"RDP": ServiceType.RDP,
|
||||
"VNC": ServiceType.VNC,
|
||||
"SMB": ServiceType.SMB,
|
||||
"Database": ServiceType.DATABASE,
|
||||
"FTP": ServiceType.FTP
|
||||
}
|
||||
return type_mapping.get(service_type_str, ServiceType.WEB_GUI)
|
||||
|
||||
|
||||
def parse_host_type(host_type_str: str) -> HostType:
|
||||
"""Convert a string to HostType enum, with fallback."""
|
||||
type_mapping = {
|
||||
"Linux": HostType.LINUX,
|
||||
"Windows": HostType.WINDOWS,
|
||||
"Windows Server": HostType.WINDOWS_SERVER,
|
||||
"Proxmox": HostType.PROXMOX,
|
||||
"ESXi": HostType.ESXI,
|
||||
"Router": HostType.ROUTER,
|
||||
"Switch": HostType.SWITCH,
|
||||
}
|
||||
return type_mapping.get(host_type_str, HostType.LINUX)
|
||||
|
||||
|
||||
def parse_vpn_type(vpn_type_str: str) -> VPNType:
|
||||
"""Convert a string to VPNType enum, with fallback."""
|
||||
type_mapping = {
|
||||
"OpenVPN": VPNType.OPENVPN,
|
||||
"WireGuard": VPNType.WIREGUARD,
|
||||
"IPSec": VPNType.IPSEC,
|
||||
}
|
||||
return type_mapping.get(vpn_type_str, VPNType.OPENVPN)
|
||||
|
||||
|
||||
def parse_host(host_data: Dict[str, Any]) -> Host:
|
||||
"""Parse a host from YAML data."""
|
||||
# Parse services
|
||||
services = []
|
||||
if 'services' in host_data:
|
||||
for service_data in host_data['services']:
|
||||
service = Service(
|
||||
name=service_data['name'],
|
||||
service_type=parse_service_type(service_data['service_type']),
|
||||
port=service_data['port']
|
||||
)
|
||||
services.append(service)
|
||||
|
||||
# Create host
|
||||
host = Host(
|
||||
name=host_data['name'],
|
||||
ip_address=host_data['ip_address'],
|
||||
host_type=parse_host_type(host_data['host_type']),
|
||||
description=host_data.get('description', ''),
|
||||
services=services
|
||||
)
|
||||
|
||||
# Parse sub-hosts (VMs) recursively
|
||||
if 'sub_hosts' in host_data:
|
||||
for subhost_data in host_data['sub_hosts']:
|
||||
subhost = parse_host(subhost_data)
|
||||
host.sub_hosts.append(subhost)
|
||||
|
||||
return host
|
||||
|
||||
|
||||
def parse_location(location_data: Dict[str, Any]) -> Location:
|
||||
"""Parse a location from YAML data."""
|
||||
# Parse hosts
|
||||
hosts = []
|
||||
if 'hosts' in location_data:
|
||||
for host_data in location_data['hosts']:
|
||||
host = parse_host(host_data)
|
||||
hosts.append(host)
|
||||
|
||||
# Create location
|
||||
location = Location(
|
||||
name=location_data['name'],
|
||||
vpn_type=parse_vpn_type(location_data['vpn_type']),
|
||||
connected=location_data.get('connected', False),
|
||||
active=location_data.get('active', False),
|
||||
vpn_config=location_data.get('vpn_config', ''),
|
||||
hosts=hosts
|
||||
)
|
||||
|
||||
return location
|
||||
|
||||
|
||||
def parse_customer(yaml_file: Path) -> Customer:
|
||||
"""Parse a customer from a YAML file."""
|
||||
with open(yaml_file, 'r') as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
# Parse customer services
|
||||
services = []
|
||||
if 'services' in data:
|
||||
for service_data in data['services']:
|
||||
service = CustomerService(
|
||||
name=service_data['name'],
|
||||
url=service_data['url'],
|
||||
service_type=service_data['service_type'],
|
||||
description=service_data.get('description', '')
|
||||
)
|
||||
services.append(service)
|
||||
|
||||
# Parse locations
|
||||
locations = []
|
||||
if 'locations' in data:
|
||||
for location_data in data['locations']:
|
||||
location = parse_location(location_data)
|
||||
locations.append(location)
|
||||
|
||||
# Create customer
|
||||
customer = Customer(
|
||||
name=data['name'],
|
||||
services=services,
|
||||
locations=locations
|
||||
)
|
||||
|
||||
return customer
|
||||
|
||||
|
||||
def load_customers() -> List[Customer]:
|
||||
"""Load customer data. Currently returns mock data for demonstration."""
|
||||
|
||||
"""Load all customers from YAML files in the config directory."""
|
||||
config_dir = ensure_config_dir()
|
||||
customers = []
|
||||
|
||||
# Customer 1: TechCorp Solutions
|
||||
techcorp = Customer(name="TechCorp Solutions")
|
||||
|
||||
# TechCorp's cloud services
|
||||
techcorp.services = [
|
||||
CustomerService("Office 365", "https://portal.office.com", "Email & Office"),
|
||||
CustomerService("Pascom Cloud PBX", "https://techcorp.pascom.cloud", "Phone System"),
|
||||
CustomerService("Salesforce CRM", "https://techcorp.salesforce.com", "CRM")
|
||||
]
|
||||
|
||||
# TechCorp's main office location
|
||||
main_office = Location(
|
||||
name="Main Office",
|
||||
vpn_type=VPNType.OPENVPN,
|
||||
connected=True,
|
||||
active=True,
|
||||
vpn_config="/etc/openvpn/techcorp-main.ovpn"
|
||||
)
|
||||
|
||||
# Proxmox hypervisor with VMs
|
||||
proxmox_host = Host(
|
||||
name="PVE-01",
|
||||
ip_address="192.168.1.10",
|
||||
host_type=HostType.PROXMOX,
|
||||
description="Main virtualization server",
|
||||
services=[
|
||||
Service("Web Interface", ServiceType.WEB_GUI, 8006),
|
||||
Service("SSH", ServiceType.SSH, 22)
|
||||
]
|
||||
)
|
||||
|
||||
# VMs running on Proxmox
|
||||
proxmox_host.sub_hosts = [
|
||||
Host(
|
||||
name="DC-01",
|
||||
ip_address="192.168.1.20",
|
||||
host_type=HostType.WINDOWS_SERVER,
|
||||
description="Domain Controller",
|
||||
services=[
|
||||
Service("RDP", ServiceType.RDP, 3389),
|
||||
Service("Admin Web", ServiceType.WEB_GUI, 8080)
|
||||
]
|
||||
),
|
||||
Host(
|
||||
name="FILE-01",
|
||||
ip_address="192.168.1.21",
|
||||
host_type=HostType.LINUX,
|
||||
description="File Server (Samba)",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("SMB Share", ServiceType.SMB, 445),
|
||||
Service("Web Panel", ServiceType.WEB_GUI, 9000)
|
||||
]
|
||||
),
|
||||
Host(
|
||||
name="DB-01",
|
||||
ip_address="192.168.1.22",
|
||||
host_type=HostType.LINUX,
|
||||
description="PostgreSQL Database",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("PostgreSQL", ServiceType.DATABASE, 5432),
|
||||
Service("pgAdmin", ServiceType.WEB_GUI, 5050)
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
# Network infrastructure
|
||||
router = Host(
|
||||
name="FW-01",
|
||||
ip_address="192.168.1.1",
|
||||
host_type=HostType.ROUTER,
|
||||
description="pfSense Firewall/Router",
|
||||
services=[
|
||||
Service("Web Interface", ServiceType.WEB_GUI, 443),
|
||||
Service("SSH", ServiceType.SSH, 22)
|
||||
]
|
||||
)
|
||||
|
||||
switch = Host(
|
||||
name="SW-01",
|
||||
ip_address="192.168.1.2",
|
||||
host_type=HostType.SWITCH,
|
||||
description="Managed Switch",
|
||||
services=[
|
||||
Service("Web Interface", ServiceType.WEB_GUI, 80),
|
||||
Service("SSH", ServiceType.SSH, 22)
|
||||
]
|
||||
)
|
||||
|
||||
main_office.hosts = [proxmox_host, router, switch]
|
||||
|
||||
# Branch office location
|
||||
branch_office = Location(
|
||||
name="Branch Office",
|
||||
vpn_type=VPNType.WIREGUARD,
|
||||
connected=False,
|
||||
active=False,
|
||||
vpn_config="/etc/wireguard/techcorp-branch.conf"
|
||||
)
|
||||
|
||||
branch_server = Host(
|
||||
name="BRANCH-01",
|
||||
ip_address="10.10.1.10",
|
||||
host_type=HostType.LINUX,
|
||||
description="Branch office server",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("File Share", ServiceType.SMB, 445),
|
||||
Service("Local Web", ServiceType.WEB_GUI, 8080)
|
||||
]
|
||||
)
|
||||
|
||||
branch_office.hosts = [branch_server]
|
||||
|
||||
techcorp.locations = [main_office, branch_office]
|
||||
customers.append(techcorp)
|
||||
|
||||
# Customer 2: MedPractice Group
|
||||
medpractice = Customer(name="MedPractice Group")
|
||||
|
||||
# MedPractice's cloud services
|
||||
medpractice.services = [
|
||||
CustomerService("Google Workspace", "https://workspace.google.com", "Email & Office"),
|
||||
CustomerService("Practice Management", "https://medpractice.emr-system.com", "EMR System"),
|
||||
CustomerService("VoIP Provider", "https://medpractice.voip.com", "Phone System")
|
||||
]
|
||||
|
||||
# Clinic location
|
||||
clinic_location = Location(
|
||||
name="Main Clinic",
|
||||
vpn_type=VPNType.WIREGUARD,
|
||||
connected=False,
|
||||
active=False,
|
||||
vpn_config="/etc/wireguard/medpractice.conf"
|
||||
)
|
||||
|
||||
# ESXi hypervisor
|
||||
esxi_host = Host(
|
||||
name="ESXi-01",
|
||||
ip_address="10.0.1.10",
|
||||
host_type=HostType.ESXI,
|
||||
description="VMware ESXi Host",
|
||||
services=[
|
||||
Service("vSphere Web", ServiceType.WEB_GUI, 443),
|
||||
Service("SSH", ServiceType.SSH, 22)
|
||||
]
|
||||
)
|
||||
|
||||
# VMs on ESXi
|
||||
esxi_host.sub_hosts = [
|
||||
Host(
|
||||
name="WIN-SRV-01",
|
||||
ip_address="10.0.1.20",
|
||||
host_type=HostType.WINDOWS_SERVER,
|
||||
description="Windows Server 2022",
|
||||
services=[
|
||||
Service("RDP", ServiceType.RDP, 3389),
|
||||
Service("IIS Web", ServiceType.WEB_GUI, 80)
|
||||
]
|
||||
),
|
||||
Host(
|
||||
name="BACKUP-01",
|
||||
ip_address="10.0.1.21",
|
||||
host_type=HostType.LINUX,
|
||||
description="Backup Server",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("Backup Web UI", ServiceType.WEB_GUI, 8080)
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
# Physical server
|
||||
physical_server = Host(
|
||||
name="PHYS-01",
|
||||
ip_address="10.0.1.50",
|
||||
host_type=HostType.LINUX,
|
||||
description="Physical Linux Server",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("Docker Portainer", ServiceType.WEB_GUI, 9000),
|
||||
Service("Nginx Proxy", ServiceType.WEB_GUI, 8080)
|
||||
]
|
||||
)
|
||||
|
||||
clinic_location.hosts = [esxi_host, physical_server]
|
||||
medpractice.locations = [clinic_location]
|
||||
customers.append(medpractice)
|
||||
|
||||
# Customer 3: Manufacturing Inc
|
||||
manufacturing = Customer(name="Manufacturing Inc")
|
||||
|
||||
# Manufacturing's cloud services
|
||||
manufacturing.services = [
|
||||
CustomerService("Microsoft 365", "https://portal.office.com", "Email & Office"),
|
||||
CustomerService("SAP Cloud", "https://manufacturing.sap.com", "ERP System")
|
||||
]
|
||||
|
||||
# Factory location
|
||||
factory_location = Location(
|
||||
name="Factory Floor",
|
||||
vpn_type=VPNType.IPSEC,
|
||||
connected=False,
|
||||
active=True,
|
||||
vpn_config="/etc/ipsec.d/manufacturing.conf"
|
||||
)
|
||||
|
||||
# Manufacturing infrastructure - simpler setup
|
||||
linux_server = Host(
|
||||
name="PROD-01",
|
||||
ip_address="172.16.1.10",
|
||||
host_type=HostType.LINUX,
|
||||
description="Production Server",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("Web Portal", ServiceType.WEB_GUI, 8443),
|
||||
Service("FTP", ServiceType.FTP, 21)
|
||||
]
|
||||
)
|
||||
|
||||
nas_server = Host(
|
||||
name="NAS-01",
|
||||
ip_address="172.16.1.20",
|
||||
host_type=HostType.LINUX,
|
||||
description="Network Attached Storage",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("Web Interface", ServiceType.WEB_GUI, 5000),
|
||||
Service("SMB Share", ServiceType.SMB, 445)
|
||||
]
|
||||
)
|
||||
|
||||
factory_location.hosts = [linux_server, nas_server]
|
||||
|
||||
# Office location
|
||||
office_location = Location(
|
||||
name="Administrative Office",
|
||||
vpn_type=VPNType.OPENVPN,
|
||||
connected=False,
|
||||
active=False,
|
||||
vpn_config="/etc/openvpn/manufacturing-office.ovpn"
|
||||
)
|
||||
|
||||
office_server = Host(
|
||||
name="OFFICE-01",
|
||||
ip_address="172.16.2.10",
|
||||
host_type=HostType.WINDOWS_SERVER,
|
||||
description="Office domain controller",
|
||||
services=[
|
||||
Service("RDP", ServiceType.RDP, 3389),
|
||||
Service("File Share", ServiceType.SMB, 445)
|
||||
]
|
||||
)
|
||||
|
||||
office_location.hosts = [office_server]
|
||||
|
||||
manufacturing.locations = [factory_location, office_location]
|
||||
customers.append(manufacturing)
|
||||
|
||||
|
||||
# Get all YAML files in the directory
|
||||
yaml_files = list(config_dir.glob("*.yaml")) + \
|
||||
list(config_dir.glob("*.yml"))
|
||||
|
||||
if not yaml_files:
|
||||
# No customer files found, initialize with examples
|
||||
print(f"No customer files found in {config_dir}")
|
||||
print("Run 'python data_loader.py --init' to create example customer files")
|
||||
return get_demo_customers()
|
||||
|
||||
# Load each customer file
|
||||
for yaml_file in yaml_files:
|
||||
try:
|
||||
customer = parse_customer(yaml_file)
|
||||
customers.append(customer)
|
||||
print(f"Loaded customer: {customer.name} from {yaml_file.name}")
|
||||
except Exception as e:
|
||||
print(f"Error loading {yaml_file}: {e}")
|
||||
|
||||
return customers
|
||||
|
||||
|
||||
def save_customers(customers: List[Customer]) -> None:
|
||||
"""Save customer data. Currently a placeholder."""
|
||||
# TODO: Implement actual persistence (JSON file, database, etc.)
|
||||
pass
|
||||
def save_customer(customer: Customer, filename: str = None) -> None:
|
||||
"""Save a customer to a YAML file."""
|
||||
config_dir = ensure_config_dir()
|
||||
|
||||
if filename is None:
|
||||
# Generate filename from customer name
|
||||
filename = customer.name.lower().replace(' ', '_') + '.yaml'
|
||||
|
||||
filepath = config_dir / filename
|
||||
|
||||
# Convert customer to dictionary
|
||||
data = {
|
||||
'name': customer.name,
|
||||
'services': [
|
||||
{
|
||||
'name': service.name,
|
||||
'url': service.url,
|
||||
'service_type': service.service_type,
|
||||
'description': service.description
|
||||
}
|
||||
for service in customer.services
|
||||
],
|
||||
'locations': []
|
||||
}
|
||||
|
||||
# Convert locations
|
||||
for location in customer.locations:
|
||||
location_data = {
|
||||
'name': location.name,
|
||||
'vpn_type': location.vpn_type.value,
|
||||
'vpn_config': location.vpn_config,
|
||||
'active': location.active,
|
||||
'connected': location.connected,
|
||||
'hosts': []
|
||||
}
|
||||
|
||||
# Convert hosts
|
||||
def convert_host(host):
|
||||
host_data = {
|
||||
'name': host.name,
|
||||
'ip_address': host.ip_address,
|
||||
'host_type': host.host_type.value,
|
||||
'description': host.description,
|
||||
'services': [
|
||||
{
|
||||
'name': service.name,
|
||||
'service_type': service.service_type.value,
|
||||
'port': service.port
|
||||
}
|
||||
for service in host.services
|
||||
]
|
||||
}
|
||||
|
||||
if host.sub_hosts:
|
||||
host_data['sub_hosts'] = [convert_host(
|
||||
subhost) for subhost in host.sub_hosts]
|
||||
|
||||
return host_data
|
||||
|
||||
for host in location.hosts:
|
||||
location_data['hosts'].append(convert_host(host))
|
||||
|
||||
data['locations'].append(location_data)
|
||||
|
||||
# Write to file
|
||||
with open(filepath, 'w') as f:
|
||||
yaml.dump(data, f, default_flow_style=False, sort_keys=False)
|
||||
|
||||
print(f"Saved customer to {filepath}")
|
||||
|
||||
|
||||
def get_demo_customers() -> List[Customer]:
|
||||
"""Return demo customers for when no config files exist."""
|
||||
# Return a minimal demo customer
|
||||
demo_customer = Customer(name="Demo Customer")
|
||||
|
||||
demo_customer.services = [
|
||||
CustomerService(
|
||||
name="Demo Portal",
|
||||
url="https://demo.example.com",
|
||||
service_type="Web Portal",
|
||||
description="Demo web portal"
|
||||
)
|
||||
]
|
||||
|
||||
demo_location = Location(
|
||||
name="Demo Location",
|
||||
vpn_type=VPNType.OPENVPN,
|
||||
connected=False,
|
||||
active=True,
|
||||
vpn_config="/etc/openvpn/demo.ovpn"
|
||||
)
|
||||
|
||||
demo_host = Host(
|
||||
name="DEMO-01",
|
||||
ip_address="10.0.0.1",
|
||||
host_type=HostType.LINUX,
|
||||
description="Demo server",
|
||||
services=[
|
||||
Service("SSH", ServiceType.SSH, 22),
|
||||
Service("Web", ServiceType.WEB_GUI, 80)
|
||||
]
|
||||
)
|
||||
|
||||
demo_location.hosts = [demo_host]
|
||||
demo_customer.locations = [demo_location]
|
||||
|
||||
return [demo_customer]
|
||||
|
||||
|
||||
def initialize_example_customers():
|
||||
"""Create example customer YAML files in the config directory."""
|
||||
config_dir = ensure_config_dir()
|
||||
|
||||
# Create TechCorp example
|
||||
techcorp_file = config_dir / "techcorp_solutions.yaml"
|
||||
if not techcorp_file.exists():
|
||||
# Read from our example file
|
||||
example_file = Path(__file__).parent / "example_customer.yaml"
|
||||
if example_file.exists():
|
||||
with open(example_file, 'r') as f:
|
||||
content = f.read()
|
||||
with open(techcorp_file, 'w') as f:
|
||||
f.write(content)
|
||||
print(f"Created example: {techcorp_file}")
|
||||
|
||||
# Create a simpler example
|
||||
simple_file = config_dir / "simple_customer.yaml"
|
||||
if not simple_file.exists():
|
||||
simple_yaml = """name: Simple Customer
|
||||
|
||||
services:
|
||||
- name: Company Website
|
||||
url: https://simple.example.com
|
||||
service_type: Web Portal
|
||||
description: Main company website
|
||||
|
||||
locations:
|
||||
- name: Main Office
|
||||
vpn_type: WireGuard
|
||||
vpn_config: /etc/wireguard/simple.conf
|
||||
active: false
|
||||
connected: false
|
||||
|
||||
hosts:
|
||||
- name: SERVER-01
|
||||
ip_address: 192.168.1.10
|
||||
host_type: Linux
|
||||
description: Main server
|
||||
services:
|
||||
- name: SSH
|
||||
service_type: SSH
|
||||
port: 22
|
||||
- name: Web Interface
|
||||
service_type: Web GUI
|
||||
port: 443
|
||||
"""
|
||||
with open(simple_file, 'w') as f:
|
||||
f.write(simple_yaml)
|
||||
print(f"Created example: {simple_file}")
|
||||
|
||||
print(f"\nExample customer files created in: {config_dir}")
|
||||
print("You can now edit these files or create new ones following the same format.")
|
||||
|
||||
|
||||
# Allow running this file directly to initialize examples
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
if len(sys.argv) > 1 and sys.argv[1] == "--init":
|
||||
initialize_example_customers()
|
||||
else:
|
||||
# Test loading
|
||||
customers = load_customers()
|
||||
for customer in customers:
|
||||
print(f"\nLoaded: {customer.name}")
|
||||
print(f" Services: {len(customer.services)}")
|
||||
print(f" Locations: {len(customer.locations)}")
|
||||
for location in customer.locations:
|
||||
print(f" - {location.name}: {len(location.hosts)} hosts")
|
||||
|
||||
Reference in New Issue
Block a user