🔮 better handler

handler/handler.c

@@ -16,56 +16,49 @@
 
 // Get command options
 char *get_commands() {
-  char *message = "===== MySecurity =====\n\n[0] Shutdown Computer\n[1] Turn "
-                  "off MySecurity\n[2] Exit\n";
+  char *message = "===== MySecurity =====\n\n[1] Shutdown Computer\n[2] Reboot"
+                  "Computer\n[0] Exit";
   char *message_ptr = message;
 
   return message_ptr;
 }
 
-// Get instance info
-int get_instance_info(char *name, char *ip) {
-  char host[256];
-  struct hostent *host_entry;
-  int hostname;
-
-  // Find the host name
-  hostname = gethostname(host, sizeof(host));
-  if (hostname < 0)
-    return -1;
+// Get option and execute a command
+int run_command(char *option) {
+#ifdef _WIN32 // Windows based operating system 
+  switch (atoi(option)) {
+  case 0:
+    return 1;
+    break;
 
-  // Find host information
-  host_entry = gethostbyname(host);
-  if (host_entry == NULL)
-    return -1;
+  case 1:
+    system("shutdown -s now");
+    break;
 
-  // Convert into IP string
-  ip = inet_ntoa(*((struct in_addr *)host_entry->h_addr_list[0]));
+  case 2:
+    system("shutdown -r");
+    break;
 
-  name = host;
+  default:
+    return -1;
+  }
+#else // Unix based operating system 
+  switch (atoi(option)) {
+  case 0:
+    return 1;
+    break;
 
-  return 0;
-}
+  case 1:
+    system("sudo shutdown -P now");
+    break;
 
-// Get option and execute a command
-int run_command(char *option, bool *connected) {
-  printf("[!] run command: %s\n", option);
-  if (option[0] == '0') {
-    printf("shutting down computer");
-#ifdef _WIN32
-    system("shutdown -s");
-#else
-    system("sudo shutdown -P");
-#endif
-    return 0;
+  case 2:
+    system("sudo reboot");
+    break;
 
-  } else if (option[0] == '1') {
-    *connected = false;
-    return 0;
-  } else if (option[0] == '2') {
-    *connected = false;
-    return 0;
+  default:
+    return -1;
   }
-
-  return -1;
+#endif
+  return 0;
 }

handler/handler.h

@@ -3,8 +3,5 @@
 // Get command options
 char *get_commands();
 
-// Get instance info
-int get_instance_info(char *name, char* ip);
-
 // Get option and execute a command
-int run_command(char *option, bool *connected);
+int run_command(char *option);

main.c

@@ -42,14 +42,18 @@ int main() {
     }
 
     // Handle user option
-    error = run_command(buffer, &connected);
+    error = run_command(buffer);
     if (error < 0) {
       error = send_socket(client, "[!] Invalid option");
       if (error < 0) {
         close_socket(client, sock);
         connected = false;
       }
     }
+    if (error == 1) {
+      close_socket(client, sock);
+      connected = false;
+    }
 
     // Clear variable value from memory
     memset(buffer, 0, sizeof(buffer));