Merge pull request #5 from isankadn/Isanka-patch-83d4c9c

fixing primary key issue

ssh_key_client/Cargo.toml

@@ -13,6 +13,7 @@ tokio = {version = "1.32.0", features = ["full"]}
 hostname = "0.3"
 pnet = "0.27"
 openssl-sys = "*"
+lazy_static = "1.4"
 
 [features]
 # Force openssl-sys to staticly link in the openssl library. Necessary when

ssh_key_client/src/main.rs

@@ -1,10 +1,21 @@
+#[macro_use]
+extern crate lazy_static;
+
 use std::env;
 use std::fs;
 use std::path::Path;
 use serde::{Serialize, Deserialize};
 use reqwest;
 use hostname;
 use pnet::datalink;
+use std::io::{self, ErrorKind};
+
+lazy_static! {
+    static ref DEBUG: bool = {
+        let debug_var = env::var("SSH_KEY_CLIENT_DEBUG").unwrap_or_default();
+        debug_var == "1" || debug_var == "1.000000000e+00"
+    };
+}
 
 #[derive(Serialize, Deserialize)]
 struct SSHKeyReport {
@@ -30,16 +41,33 @@ fn get_primary_ip() -> Option<String> {
 
 fn read_ssh_keys(path: &Path) -> Option<Vec<String>> {
     if let Ok(content) = fs::read_to_string(path) {
+        if *DEBUG {
+            println!("Reading content: {}", content);
+        }
         Some(content.lines().map(|s| s.to_string()).collect())
     } else {
         None
     }
+
 }
 
-fn get_vm_uuid() -> Result<String, std::io::Error> {
-    fs::read_to_string("/sys/class/dmi/id/product_uuid")
+fn get_vm_uuid() -> Result<String, io::Error> {
+    println!("debug enabled: {}", DEBUG.to_string());
+
+    if *DEBUG {
+        println!("Reading VM UUID from /sys/class/dmi/id/product_uuid");
+    }
+    match fs::read_to_string("/sys/class/dmi/id/product_uuid") {
+        Ok(uuid) => Ok(uuid),
+        Err(e) if e.kind() == ErrorKind::PermissionDenied => {
+            eprintln!("Permission denied. Please run the program with sudo.");
+            std::process::exit(1);
+        }
+        Err(e) => Err(e),
+    }
 }
 
+
 async fn send_to_server(report: SSHKeyReport, server_url: &str) -> Result<(), reqwest::Error> {
     let client = reqwest::Client::new();
     client.post(server_url).json(&report).send().await?;
@@ -49,6 +77,9 @@ async fn send_to_server(report: SSHKeyReport, server_url: &str) -> Result<(), re
 // This function reads the /etc/passwd file and returns a list of home directories
 fn get_user_home_dirs() -> Vec<String> {
     let content = fs::read_to_string("/etc/passwd").unwrap_or_default();
+    if *DEBUG {
+        println!("Reading content: {:?}", content);
+    }
     content.lines()
         .filter_map(|line| {
             let parts: Vec<&str> = line.split(':').collect();
@@ -59,6 +90,7 @@ fn get_user_home_dirs() -> Vec<String> {
             }
         })
         .collect()
+
 }
 
 #[tokio::main]
@@ -82,6 +114,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     for home_dir in get_user_home_dirs() {
         let key_path = Path::new(&home_dir).join(".ssh/authorized_keys");
+        if *DEBUG {
+            println!("Reading SSH keys from: {:?}", key_path);
+        }     
         if let Some(keys) = read_ssh_keys(&key_path) {
             let report = SSHKeyReport {
                 vm_name: vm_name.clone(),

ssh_key_server/src/main.rs

@@ -14,6 +14,7 @@ use serde_yaml;
 use base64;
 use rocket::http::{Cookie, Cookies};
 use rocket::response::Redirect;
+use rocket_contrib::serve::StaticFiles;
 
 
 struct AuthenticatedUser(String);
@@ -145,8 +146,20 @@ impl<'a, 'r> FromRequest<'a, 'r> for MaybeAuthenticatedUser {
 
 #[post("/", data = "<report>")]
 fn receive_keys(report: Json<SSHKeyReport>, storage: State<'_, KeyStorage>) -> &'static str {
-    let mut db = storage.lock().expect("Failed to lock storage.");
-    db.insert(report.vm_uuid.clone(), report.clone()); // Use vm_uuid as the key
+    let mut db = match storage.lock() {
+        Ok(guard) => guard,
+        Err(_) => return "Failed to lock storage.",
+    };
+
+    // Iterate over each key in the report's keys vector
+    for key in report.keys.iter() {
+        // Create a composite key using vm_uuid and the individual key
+        let composite_key = format!("{}-{}", report.vm_uuid, key);
+
+        // Insert the data into the storage using the composite key
+        db.insert(composite_key, report.clone()); 
+    }
+
     "Received"
 }
 
@@ -217,7 +230,8 @@ fn rocket() -> rocket::Rocket {
     rocket::ignite()
         .attach(Template::fairing()) // Attach the template fairing
         .manage(Mutex::new(HashMap::<String, SSHKeyReport>::new()))
-        .mount("/", routes![receive_keys, list_keys, login_page, login_submit])
+        .mount("/static", StaticFiles::from("static"))
+        .mount("/", routes![receive_keys, list_keys, login_page, login_submit])        
         .register(catchers![not_found])
 }
 

ssh_key_server/templates/index.hbs

@@ -11,6 +11,12 @@
     <link href="https://unpkg.com/[email protected]/dist/tailwind.min.css" rel=" stylesheet">
     <!--Replace with your tailwind.css once created-->
 
+    <link rel="apple-touch-icon" sizes="76x76" href="/static/apple-touch-icon.png">
+    <link rel="icon" type="image/png" sizes="32x32" href="/static/favicon-32x32.png">
+    <link rel="icon" type="image/png" sizes="16x16" href="/static/favicon-16x16.png">
+    <link rel="mask-icon" href="/static/safari-pinned-tab.svg" color="#5bbad5">
+    <meta name="msapplication-TileColor" content="#da532c">
+    <meta name="theme-color" content="#ffffff">
 
     <!--Regular Datatables CSS-->
     <link href="https://cdn.datatables.net/1.10.19/css/jquery.dataTables.min.css" rel="stylesheet">
@@ -125,7 +131,6 @@
         <!--Card-->
         <div id='recipients' class="p-8 mt-6 lg:mt-0 rounded shadow bg-white">
 
-
             <table id="keytable" class="stripe hover" style="width:100%; padding-top: 1em;  padding-bottom: 1em;">
                 <thead>
                     <tr>
@@ -187,12 +192,14 @@
     <script>
         $(document).ready(function () {
 
-            var table = $('#keytable').DataTable({
+            var table = $('#keytable').DataTable({                
                 responsive: true
             })
                 .columns.adjust()
                 .responsive.recalc();
-            $('#customSearchBox').on('keyup', function () {
+            $('.dataTables_filter input').attr('id', 'searchBox')
+            $('#searchBox').on('keyup', function () {
+                console.log(this.value);
                 table.search(this.value).draw();
             });
             $('.show-key-btn').click(function () {

ssh_key_server/templates/login.html.hbs

@@ -4,8 +4,15 @@
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Login</title>
+    <title>SSHKeySight | Login</title>
     <link href="https://unpkg.com/[email protected]/dist/tailwind.min.css" rel="stylesheet">
+    <link rel="apple-touch-icon" sizes="76x76" href="/static/apple-touch-icon.png">
+    <link rel="icon" type="image/png" sizes="32x32" href="/static/favicon-32x32.png">
+    <link rel="icon" type="image/png" sizes="16x16" href="/static/favicon-16x16.png">
+    <link rel="manifest" href="/site.webmanifest">
+    <link rel="mask-icon" href="/safari-pinned-tab.svg" color="#5bbad5">
+    <meta name="msapplication-TileColor" content="#da532c">
+    <meta name="theme-color" content="#ffffff">
 </head>
 
 <body class="bg-gray-100 h-screen flex justify-center items-center">