diff --git a/bin/NETRAVE b/bin/NETRAVE index 89e5994..9bddb4e 100644 --- a/bin/NETRAVE +++ b/bin/NETRAVE @@ -87,4 +87,4 @@ end # End of the program # wait for the alert_queue_manager to block before we exit. -@alert_queue_manager.join_worker +@alert_queue_manager.shutdown diff --git a/lib/Gemfile b/lib/Gemfile index 656bfe1..4a0537e 100644 --- a/lib/Gemfile +++ b/lib/Gemfile @@ -31,12 +31,12 @@ gem 'securerandom', '~> 0.2.2' gem 'dotenv', '~> 2.8' -gem "tracer", "~> 0.2.2" +gem 'tracer', '~> 0.2.2' -gem "flay", "~> 2.13" +gem 'flay', '~> 2.13' -gem "pcaprub", "~> 0.13.1" +gem 'pcaprub', '~> 0.13.1' -gem "packetfu", "~> 2.0" +gem 'packetfu', '~> 2.0' -gem "sudo", "~> 0.2.0" +gem 'sudo', '~> 0.2.0' diff --git a/lib/utils/alert_manager.rb b/lib/utils/alert_manager.rb index 634c15b..ca67162 100644 --- a/lib/utils/alert_manager.rb +++ b/lib/utils/alert_manager.rb @@ -3,7 +3,7 @@ # Class for creating and displaying alerts in the Curses TUI. This class also manages a little bit of concurrency # We use mutex for sync so we don't break Curses, as Curses isn't thread safe class Alert - attr_reader :message, :severity + attr_reader :message, :severity, :alert_window def initialize(message, severity) @message = message @@ -11,7 +11,7 @@ class Alert @curses_mutex = Mutex.new end - def display + def display # rubocop:disable Metrics/AbcSize, Metrics/MethodLength @curses_mutex.synchronize do # Initialize color pairs Curses.start_color @@ -20,33 +20,30 @@ class Alert Curses.init_pair(3, Curses::COLOR_YELLOW, Curses::COLOR_BLACK) # Warning # Create a new window for the alert at the bottom of the screen - alert_window = Curses::Window.new(1, Curses.cols, Curses.lines - 1, 0) + @alert_window = Curses::Window.new(1, Curses.cols, Curses.lines - 1, 0) # Set the color attribute based on the severity of the alert case @severity when :info - alert_window.attron(Curses.color_pair(1) | Curses::A_NORMAL) # Blue color + @alert_window.attron(Curses.color_pair(1) | Curses::A_NORMAL) # Blue color when :warning - alert_window.attron(Curses.color_pair(3) | Curses::A_NORMAL) # Yellow color + @alert_window.attron(Curses.color_pair(3) | Curses::A_NORMAL) # Yellow color when :error - alert_window.attron(Curses.color_pair(2) | Curses::A_NORMAL) # Red color + @alert_window.attron(Curses.color_pair(2) | Curses::A_NORMAL) # Red color end # Add the message to the window and refresh it to display the message - alert_window.addstr(@message) - alert_window.refresh + @alert_window.addstr(@message) + @alert_window.refresh + end + end - # Create a new thread to handle the delay and clearing of the alert - # This is done in a separate thread to prevent the entire program from - # pausing while the alert is displayed - Thread.new do - sleep(5) # Pause for 5 seconds - - # Clear the alert - alert_window.clear - alert_window.refresh - alert_window.close - end + def clear + @curses_mutex.synchronize do + # Clear the alert + @alert_window.clear + @alert_window.refresh + @alert_window.close end end end diff --git a/lib/utils/alert_queue_manager.rb b/lib/utils/alert_queue_manager.rb index dd093d8..a3e2caf 100644 --- a/lib/utils/alert_queue_manager.rb +++ b/lib/utils/alert_queue_manager.rb @@ -3,24 +3,43 @@ require_relative 'ring_buffer' # Class for managing the queues for alerts class AlertQueueManager - def initialize(logger, size = 2 * 1024 * 1024) + SHUTDOWN_SIGNAL = 'shutdown' + + def initialize(logger, size = 2 * 1024 * 1024) # rubocop:disable Metrics/MethodLength @loggman = logger @queue = RingBuffer.new(@loggman, size) + @shutdown = false # Start a thread that continuously checks the queue and displays alerts @worker_thread = Thread.new do loop do + break if @shutdown && @queue.empty? + + if @queue.empty? + sleep(0.1) # Sleep for 100 milliseconds + next + end + alert = @queue.pop # This will block until there's an alert in the queue - alert&.display + next if alert.nil? + + break if alert.message == SHUTDOWN_SIGNAL + + alert.display + sleep(4.5) + alert.clear end end end def enqueue_alert(alert) + return if @shutdown + @queue.push(alert) end - def join_worker - @worker_thread.join + def shutdown + @shutdown = true + enqueue_alert(Alert.new(SHUTDOWN_SIGNAL, :info)) end end diff --git a/lib/utils/first_run_init.rb b/lib/utils/first_run_init.rb index 5e2791a..958f162 100644 --- a/lib/utils/first_run_init.rb +++ b/lib/utils/first_run_init.rb @@ -13,10 +13,12 @@ class FirstRunInit include Utilities include Curses - def initialize(loggman, alert_queue_manager, db_manager = nil) - @loggman = loggman + def initialize(logger, alert_queue_manager, db_manager = nil) + @loggman = logger @db_manager = db_manager @alert_queue_manager = alert_queue_manager + @info_gatherer = SystemInformationGather.new(@db_manager, @loggman) + Dotenv.load end def run