From 11ad66517358f8b5475486efb52c5171d54bfcb1 Mon Sep 17 00:00:00 2001
From: Aleksandr Tcitlionok <803797+terghalin@users.noreply.github.com>
Date: Mon, 9 Dec 2024 03:21:25 +0000
Subject: [PATCH] update(vm): use token ENV and skip dup

---
 app/routes/vm.py | 38 ++++++++++++++++++++++++++------------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/app/routes/vm.py b/app/routes/vm.py
index 007000b..335d161 100644
--- a/app/routes/vm.py
+++ b/app/routes/vm.py
@@ -2,9 +2,12 @@ from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 from database import insert_virtual_machine, fetch_all
 import requests
+import os
 from datetime import datetime, timezone
+import logging
 
 router = APIRouter()
+logger = logging.getLogger(__name__)
 
 class VirtualMachine(BaseModel):
     name: str
@@ -14,15 +17,16 @@ class VirtualMachine(BaseModel):
     storage: str
     type: str
 
-class HetznerToken(BaseModel):
-    api_token: str
-
 @router.get("/vm/data")
 def get_vm_data():
     return {"virtual_machines": fetch_all("virtual_machines")}
 
 @router.post("/vm/data")
 def add_vm_data(vm: VirtualMachine):
+    existing_vms = fetch_all("virtual_machines")
+    if any(existing_vm[1] == vm.name for existing_vm in existing_vms):
+        return {"message": f"Virtual machine '{vm.name}' already exists."}
+    
     insert_virtual_machine(
         name=vm.name,
         location=vm.location,
@@ -34,14 +38,18 @@ def add_vm_data(vm: VirtualMachine):
     return {"message": f"Virtual machine '{vm.name}' added successfully."}
 
 @router.post("/vm/import-hetzner")
-def import_hetzner_vms(token: HetznerToken):
+def import_hetzner_vms(api_token: str = None):
     """
     Import virtual machines from Hetzner Cloud.
     """
+    # Check for Hetzner token
+    api_token = api_token or os.getenv("HETZNER_TOKEN")
+    if not api_token:
+        logger.warning("HETZNER_TOKEN is not set. Unable to fetch Hetzner servers.")
+        return {"message": "HETZNER_TOKEN is not set. Please set it to fetch servers from Hetzner."}
+
     hetzner_url = "https://api.hetzner.cloud/v1/servers"
-    headers = {
-        "Authorization": f"Bearer {token.api_token}"
-    }
+    headers = {"Authorization": f"Bearer {api_token}"}
 
     try:
         response = requests.get(hetzner_url, headers=headers)
@@ -53,17 +61,22 @@ def import_hetzner_vms(token: HetznerToken):
     if not servers:
         return {"message": "No servers found on Hetzner Cloud."}
 
+    # Fetch existing VMs from the database
+    existing_vms = {vm[1] for vm in fetch_all("virtual_machines")}
+
+    imported_count = 0
     for server in servers:
         try:
             name = server["name"]
+            if name in existing_vms:
+                logger.info(f"Skipping existing virtual machine: {name}")
+                continue
+
             location = server["datacenter"]["location"]["name"]
             cpu = server["server_type"]["cores"]
             memory = f"{server['server_type']['memory']} GB"
             storage = f"{server['server_type']['disk']} GB"
             vm_type = server["server_type"]["name"]
-            created_at = datetime.fromisoformat(server["created"].rstrip("Z")).replace(tzinfo=timezone.utc)
-            now = datetime.now(timezone.utc)
-            time_on_duty = round((now - created_at).total_seconds() / 3600, 2)  # Convert to hours
 
             insert_virtual_machine(
                 name=name,
@@ -73,7 +86,8 @@ def import_hetzner_vms(token: HetznerToken):
                 storage=storage,
                 vm_type=vm_type
             )
+            imported_count += 1
         except Exception as e:
-            print(f"Failed to import server {server['name']}: {e}")
+            logger.error(f"Failed to import server {server['name']}: {e}")
 
-    return {"message": f"Imported {len(servers)} virtual machines from Hetzner."}
+    return {"message": f"Imported {imported_count} new virtual machines from Hetzner."}