From b3dbd0f07c9b4b7266b906433ac53a97a6945ae5 Mon Sep 17 00:00:00 2001 From: VetheonGames Date: Thu, 29 Jun 2023 18:27:08 -0600 Subject: [PATCH] Refactor and enhance database connection and data validation - Refactored the entry-point to bypass first run setup if valid configuration information already exists. - Enhanced data validation before inserting into the database. - Ensured correct data placement in the appropriate tables and columns in the database. - Added logging for database connection attempts and results. - Fixed a bug in the speed conversion from Gbps to Mbps. --- bin/NETRAVE | 50 ++++++++++++++++--------- lib/utils/database_manager.rb | 52 ++++++++++++++++++++------ lib/utils/first_run_init.rb | 18 ++++++--- lib/utils/system_information_gather.rb | 9 +++-- lib/utils/utilities.rb | 14 +++++-- 5 files changed, 103 insertions(+), 40 deletions(-) diff --git a/bin/NETRAVE b/bin/NETRAVE index 4a8099d..f989160 100644 --- a/bin/NETRAVE +++ b/bin/NETRAVE @@ -11,7 +11,7 @@ require_relative '../lib/utils/logg_man' include Utilities # rubocop:disable Style/MixinUsage # binding.b(do: 'irb') -loggman = LoggMan.new +@loggman = LoggMan.new # Create .env file if it doesn't exist File.open('.env', 'w') {} unless File.exist?('.env') @@ -20,7 +20,7 @@ File.open('.env', 'w') {} unless File.exist?('.env') Dotenv.load # Initialize DatabaseManager -db_manager = DatabaseManager.new +db_manager = DatabaseManager.new(@loggman) # Get database details from environment variables db_details = { @@ -30,34 +30,50 @@ db_details = { database: ENV['DB_DATABASE'] } +# Decrypt password +dec_pass = decrypt_string_chacha20(db_details[:password], db_details[:key]) + # If any of the necessary details are missing, run the first run setup if db_details.values.any?(&:nil?) - loggman.log_warn('Missing or incomplete configuration. Running first run setup.') - first_run_init = FirstRunInit.new(db_manager) + @loggman.log_warn('Missing or incomplete configuration. Running first run setup.') + first_run_init = FirstRunInit.new(@loggman, db_manager) first_run_init.run # Reload environment variables after first run setup Dotenv.load - username = ENV['DB_USERNAME'] - password = ENV['DB_PASSWORD'] - key = ENV['DB_SECRET_KEY'] - database = ENV['DB_DATABASE'] + db_details = { + username: ENV['DB_USERNAME'], + password: ENV['DB_PASSWORD'], + key: ENV['DB_SECRET_KEY'], + database: ENV['DB_DATABASE'] + } + # Decrypt password again after potentially updating config + dec_pass = decrypt_string_chacha20(db_details[:password], db_details[:key]) end # Test connection -dec_pass = decrypt_string_chacha20(password, key) -unless db_manager.test_db_connection(username, dec_pass, database) - loggman.log_warn('Failed to connect to the database with existing configuration. Please re-enter your details.') - first_run_init = FirstRunInit.new(db_manager) +unless db_manager.test_db_connection(db_details[:username], dec_pass, db_details[:database]) + @loggman.log_warn('Failed to connect to the database with existing configuration. Please re-enter your details.') + first_run_init = FirstRunInit.new(@loggman, db_manager) first_run_init.run + # Reload environment variables after potentially updating config + Dotenv.load + db_details = { + username: ENV['DB_USERNAME'], + password: ENV['DB_PASSWORD'], + key: ENV['DB_SECRET_KEY'], + database: ENV['DB_DATABASE'] + } + # Decrypt password again after potentially updating config + dec_pass = decrypt_string_chacha20(db_details[:password], db_details[:key]) end # Test connection again after potentially updating config -if db_manager.test_db_connection(username, dec_pass, database) - loggman.log_info('Successfully connected to the database.') +if db_manager.test_db_connection(db_details[:username], dec_pass, db_details[:database]) + @loggman.log_info('Successfully connected to the database.') else - loggman.log_error('Failed to connect to the database. Please check your configuration.') + @loggman.log_error('Failed to connect to the database. Please check your configuration.') exit 1 end -puts 'Program successfully ran with no errors' -# TODO: Add the rest of your application logic here +@loggman.log_warn('Program successfully ran with no errors') +# TODO: Add the rest of application logic here diff --git a/lib/utils/database_manager.rb b/lib/utils/database_manager.rb index 9957159..dc1612c 100644 --- a/lib/utils/database_manager.rb +++ b/lib/utils/database_manager.rb @@ -4,20 +4,18 @@ require 'sequel' require 'mysql2' require_relative 'system_information_gather' require_relative '../utils/utilities' -require_relative 'logg_man' # database manager class DatabaseManager include Utilities - def initialize + def initialize(logger) @db = nil - @loggman = LoggMan.new + @loggman = logger end def test_db_connection(username, password, database) # rubocop:disable Metrics/MethodLength - loggman = LoggMan.new - loggman.log_info('Attempting to connect to the database...') + @loggman.log_info('Attempting to connect to the database...') display_alert('Attempting to connect to the database...', :info) # Create the connection string @@ -25,11 +23,11 @@ class DatabaseManager @db = Sequel.connect(connection_string) # Try a simple query to test the connection @db.run 'SELECT 1' - loggman.log_info('Successfully connected to the database.') + @loggman.log_info('Successfully connected to the database.') display_alert('Successfully connected to the database.', :info) true rescue Sequel::DatabaseConnectionError => e - loggman.log_error("Failed to connect to the database: #{e.message}") + @loggman.log_error("Failed to connect to the database: #{e.message}") display_alert('Failed to connect to the database!', :error) false end @@ -63,8 +61,25 @@ class DatabaseManager end end - def store_system_info(system_info) - @db[:system_info].insert(system_info) + def store_system_info(system_info) # rubocop:disable Metrics/MethodLength + # Check if the system_info already exists in the database + @loggman.log_info('Checking if info exists in the Database...') + + existing_system_info = @db[:system_info].where(uplink_speed: system_info[:uplink_speed], + downlink_speed: system_info[:downlink_speed], + total_bandwidth: system_info[:total_bandwidth]).first + + if existing_system_info + # If it exists, update it + @loggman.log_info('Info already exists. Updating instead of adding more data to the table...') + + @db[:system_info].where(id: existing_system_info[:id]).update(system_info) + else + # If it doesn't exist, insert it + @loggman.log_info('Info does not exist already, inserting it...') + + @db[:system_info].insert(system_info) + end end def create_services_table @@ -75,9 +90,24 @@ class DatabaseManager end end - def store_services(services) + def store_services(services) # rubocop:disable Metrics/MethodLength services.each do |service| - @db[:services].insert(service_name: service, status: true) + # Check if the service already exists in the database + @loggman.log_info('Checking if info exists in the Database...') + + existing_service = @db[:services].where(service_name: service).first + + if existing_service + # If it exists, update it + @loggman.log_info('Info already exists, updating instead of adding more data to the table...') + + @db[:services].where(id: existing_service[:id]).update(service_name: service, status: true) + else + # If it doesn't exist, insert it + @loggman.log_info('Info does not exist already, inserting it...') + + @db[:services].insert(service_name: service, status: true) + end end end end diff --git a/lib/utils/first_run_init.rb b/lib/utils/first_run_init.rb index c5c0953..87755c0 100644 --- a/lib/utils/first_run_init.rb +++ b/lib/utils/first_run_init.rb @@ -6,17 +6,16 @@ require 'dotenv' require_relative 'database_manager' require_relative 'system_information_gather' require_relative 'utilities' -require_relative 'logg_man' # first run class class FirstRunInit include Utilities include Curses - def initialize(db_manager = nil) - @db_manager = db_manager || DatabaseManager.new - @info_gatherer = SystemInformationGather.new(@db_manager) - @loggman = LoggMan.new + def initialize(logger, db_manager = nil) + @db_manager = db_manager || DatabaseManager.new(logger) + @info_gatherer = SystemInformationGather.new(@db_manager, logger) + @loggman = logger Dotenv.load end @@ -57,33 +56,40 @@ class FirstRunInit end def ask_for_db_details # rubocop:disable Metrics/MethodLength, Metrics/AbcSize + @loggman.log_info('Asking for Database details...') Curses.clear Curses.setpos(1, 0) Curses.addstr('Please enter your database username: ') Curses.refresh username = DCI.catch_input(true) + @loggman.log_info('Database Username entered!') Curses.setpos(2, 0) Curses.addstr('Please enter your database password: ') Curses.refresh Curses.noecho password = DCI.catch_input(false) + @loggman.log_info('Database Password Stored Securely!') Curses.echo Curses.setpos(3, 0) Curses.addstr('Please enter your database name: ') Curses.refresh database = DCI.catch_input(true) + @loggman.log_info('Database Name entered!') # Generate a secret key key = generate_key + @loggman.log_info('Secret Key Generated!') # Encrypt the password encrypted_password = encrypt_string_chacha20(password, key) + @loggman.log_info('Password Encrypted!') db_details = { username:, password: encrypted_password, key:, database: } write_db_details_to_config_file(db_details) + @loggman.log_info('Wiriting Database details to a file!') end def write_db_details_to_config_file(db_details) @@ -95,8 +101,10 @@ class FirstRunInit file.puts %(DB_DATABASE="#{db_details[:database]}") end + @loggman.log_info('Database details saved! Reloading environment...') # Load the .env file using dotenv Dotenv.load + @loggman.log_info('Environment restarted!') rescue StandardError => e @loggman.log_error("Failed to write to .env file: #{e.message}") end diff --git a/lib/utils/system_information_gather.rb b/lib/utils/system_information_gather.rb index 9e25642..8aae1aa 100644 --- a/lib/utils/system_information_gather.rb +++ b/lib/utils/system_information_gather.rb @@ -10,8 +10,9 @@ require 'dynamic_curses_input' class SystemInformationGather include Utilities - def initialize(db_manager) + def initialize(db_manager, logger) @db_manager = db_manager + @loggman = logger end def gather_system_info # rubocop:disable Metrics/MethodLength @@ -43,7 +44,7 @@ class SystemInformationGather Curses.refresh Curses.addstr('Uplink Speed: ') speed = DCI.catch_input(true) - return speed.end_with?('gbps') ? convert_speed_to_mbps(speed) : speed.to_i if valid_speed?(speed) + return convert_speed_to_mbps(speed) if valid_speed?(speed) Curses.setpos(5, 0) Curses.addstr("Whoops! That didn't appear to be a valid speed. Please try again!") @@ -60,7 +61,7 @@ class SystemInformationGather Curses.refresh Curses.addstr('Downlink Speed: ') speed = DCI.catch_input(true) - return speed.end_with?('gbps') ? convert_speed_to_mbps(speed) : speed.to_i if valid_speed?(speed) + return convert_speed_to_mbps(speed) if valid_speed?(speed) Curses.setpos(5, 0) Curses.addstr("Whoops! That didn't appear to be a valid speed. Please try again!") @@ -69,7 +70,7 @@ class SystemInformationGather end def valid_speed?(speed) - speed.to_i.positive? + speed.to_i.positive? && speed.match?(/\A\d+(gbps|mbps)\z/i) end def ask_for_services # rubocop:disable Metrics/MethodLength diff --git a/lib/utils/utilities.rb b/lib/utils/utilities.rb index 7782348..33fb94a 100644 --- a/lib/utils/utilities.rb +++ b/lib/utils/utilities.rb @@ -4,15 +4,23 @@ require 'securerandom' require 'digest' require 'base64' require 'openssl' -require_relative 'logg_man' # Utiltiies Module module Utilities # Converts speed from Gbps to Mbps if necessary def convert_speed_to_mbps(speed) - return nil unless speed.is_a?(String) && speed.match?(/\A\d+(gbps|mbps)\z/i) + return nil unless speed.is_a?(String) && speed.downcase.match?(/\A\d+(gbps|mbps)\z/i) - speed.end_with?('gbps') ? speed.to_i * 1000 : speed.to_i + # Extract the numeric part and the unit from the speed + numeric_speed, unit = speed.downcase.match(/(\d+)(gbps|mbps)/i).captures + + # Convert the numeric part to an integer + numeric_speed = numeric_speed.to_i + + # If the unit is 'gbps', multiply the numeric part by 1000 + numeric_speed *= 1000 if unit == 'gbps' + + numeric_speed end # Converts an array of services into a hash