lanet 0.2.0 → 0.3.0

data/lib/lanet/file_transfer.rb

@@ -0,0 +1,308 @@
+# frozen_string_literal: true
+
+require "digest"
+require "fileutils"
+require "tempfile"
+require "zlib"
+require "securerandom"
+require "base64"
+require "json"
+require "socket"
+require "timeout"
+
+module Lanet
+  class FileTransfer
+    # Constants
+    # Use smaller chunks in tests to avoid "Message too long" errors
+    CHUNK_SIZE = if ENV["LANET_TEST_CHUNK_SIZE"]
+                   ENV["LANET_TEST_CHUNK_SIZE"].to_i
+                 elsif ENV["RACK_ENV"] == "test"
+                   8192 # 8KB in test environment
+                 else
+                   65_536 # 64KB in production
+                 end
+
+    MAX_RETRIES = 3
+    TIMEOUT = ENV["RACK_ENV"] == "test" ? 2 : 10 # Seconds
+
+    # Message types
+    FILE_HEADER = "FH"  # File metadata
+    FILE_CHUNK  = "FC"  # File data chunk
+    FILE_END    = "FE"  # End of transfer
+    FILE_ACK    = "FA"  # Acknowledgment
+    FILE_ERROR  = "FR"  # Error message
+
+    # Custom error class
+    class Error < StandardError; end
+
+    # Attributes for tracking progress
+    attr_reader :progress, :file_size, :transferred_bytes
+
+    ### Initialization
+    def initialize(port = nil)
+      @port = port || 5001 # Default port for file transfers
+      @progress = 0.0
+      @file_size = 0
+      @transferred_bytes = 0
+      @sender = Lanet::Sender.new(@port) # Assumes Lanet::Sender is defined elsewhere
+      @cancellation_requested = false
+    end
+
+    ### Send File Method
+    def send_file(target_ip, file_path, encryption_key = nil, private_key = nil, progress_callback = nil)
+      # Validate file
+      unless File.exist?(file_path) && File.file?(file_path)
+        raise Error, "File not found or is not a regular file: #{file_path}"
+      end
+
+      # Initialize transfer state
+      @file_size = File.size(file_path)
+      @transferred_bytes = 0
+      @progress = 0.0
+      @cancellation_requested = false
+      transfer_id = SecureRandom.uuid
+      chunk_index = 0
+
+      receiver = nil
+
+      begin
+        # Send file header
+        file_name = File.basename(file_path)
+        file_checksum = calculate_file_checksum(file_path)
+        header_data = {
+          id: transfer_id,
+          name: file_name,
+          size: @file_size,
+          checksum: file_checksum,
+          timestamp: Time.now.to_i
+        }.to_json
+        header_message = Lanet::Encryptor.prepare_message("#{FILE_HEADER}#{header_data}", encryption_key, private_key)
+        @sender.send_to(target_ip, header_message)
+
+        # Wait for initial ACK
+        receiver = UDPSocket.new
+        receiver.bind("0.0.0.0", @port)
+        wait_for_ack(receiver, target_ip, transfer_id, encryption_key, "initial")
+
+        # Send file chunks
+        File.open(file_path, "rb") do |file|
+          until file.eof? || @cancellation_requested
+            chunk = file.read(CHUNK_SIZE)
+            chunk_data = {
+              id: transfer_id,
+              index: chunk_index,
+              data: Base64.strict_encode64(chunk)
+            }.to_json
+            chunk_message = Lanet::Encryptor.prepare_message("#{FILE_CHUNK}#{chunk_data}", encryption_key, private_key)
+            @sender.send_to(target_ip, chunk_message)
+
+            chunk_index += 1
+            @transferred_bytes += chunk.bytesize
+            @progress = (@transferred_bytes.to_f / @file_size * 100).round(2)
+            progress_callback&.call(@progress, @transferred_bytes, @file_size)
+
+            sleep(0.01) # Prevent overwhelming the receiver
+          end
+        end
+
+        # Send end marker and wait for final ACK
+        unless @cancellation_requested
+          end_data = { id: transfer_id, total_chunks: chunk_index }.to_json
+          end_message = Lanet::Encryptor.prepare_message("#{FILE_END}#{end_data}", encryption_key, private_key)
+          @sender.send_to(target_ip, end_message)
+          wait_for_ack(receiver, target_ip, transfer_id, encryption_key, "final")
+          true # Transfer successful
+        end
+      rescue StandardError => e
+        send_error(target_ip, transfer_id, e.message, encryption_key, private_key)
+        raise Error, "File transfer failed: #{e.message}"
+      ensure
+        receiver&.close
+      end
+    end
+
+    ### Receive File Method
+    def receive_file(output_dir, encryption_key = nil, public_key = nil, progress_callback = nil)
+      FileUtils.mkdir_p(output_dir) unless Dir.exist?(output_dir)
+      receiver = UDPSocket.new
+      receiver.bind("0.0.0.0", @port)
+      active_transfers = {}
+
+      begin
+        loop do
+          data, addr = receiver.recvfrom(65_536) # Large buffer for chunks
+          sender_ip = addr[3]
+          result = Lanet::Encryptor.process_message(data, encryption_key, public_key)
+          next unless result[:content]&.length&.> 2
+
+          message_type = result[:content][0..1]
+          message_data = result[:content][2..]
+
+          case message_type
+          when FILE_HEADER
+            handle_file_header(sender_ip, message_data, active_transfers, encryption_key, progress_callback)
+          when FILE_CHUNK
+            handle_file_chunk(sender_ip, message_data, active_transfers, progress_callback, encryption_key)
+          when FILE_END
+            handle_file_end(sender_ip, message_data, active_transfers, output_dir, encryption_key, progress_callback)
+          when FILE_ERROR
+            handle_file_error(sender_ip, message_data, active_transfers, progress_callback)
+          end
+        end
+      rescue Interrupt
+        puts "\nFile receiver stopped."
+      ensure
+        cleanup_transfers(active_transfers)
+        receiver.close
+      end
+    end
+
+    ### Cancel Transfer
+    def cancel_transfer
+      @cancellation_requested = true
+    end
+
+    private
+
+    ### Helper Methods
+
+    def calculate_file_checksum(file_path)
+      Digest::SHA256.file(file_path).hexdigest
+    end
+
+    def send_error(target_ip, transfer_id, message, encryption_key, private_key = nil)
+      error_data = { id: transfer_id, message: message, timestamp: Time.now.to_i }.to_json
+      error_message = Lanet::Encryptor.prepare_message("#{FILE_ERROR}#{error_data}", encryption_key, private_key)
+      @sender.send_to(target_ip, error_message)
+    end
+
+    def wait_for_ack(receiver, target_ip, transfer_id, encryption_key, context)
+      Timeout.timeout(TIMEOUT) do
+        data, addr = receiver.recvfrom(1024)
+        sender_ip = addr[3]
+        if sender_ip == target_ip
+          result = Lanet::Encryptor.process_message(data, encryption_key)
+          return if result[:content]&.start_with?(FILE_ACK) && result[:content][2..] == transfer_id
+
+          # Valid ACK received
+
+          raise Error, "Invalid #{context} ACK received: #{result[:content]}"
+
+        end
+      end
+    rescue Timeout::Error
+      raise Error, "Timeout waiting for #{context} transfer acknowledgment"
+    end
+
+    def handle_file_header(sender_ip, message_data, active_transfers, encryption_key, callback)
+      header = JSON.parse(message_data)
+      transfer_id = header["id"]
+      active_transfers[transfer_id] = {
+        sender_ip: sender_ip,
+        file_name: header["name"],
+        file_size: header["size"],
+        expected_checksum: header["checksum"],
+        temp_file: Tempfile.new([File.basename(header["name"], ".*"), File.extname(header["name"])]),
+        chunks_received: 0,
+        timestamp: Time.now
+      }
+      ack_message = Lanet::Encryptor.prepare_message("#{FILE_ACK}#{transfer_id}", encryption_key)
+      @sender.send_to(sender_ip, ack_message)
+      callback&.call(:start, {
+                       transfer_id: transfer_id,
+                       sender_ip: sender_ip,
+                       file_name: header["name"],
+                       file_size: header["size"]
+                     })
+    rescue JSON::ParserError => e
+      send_error(sender_ip, "unknown", "Invalid header format: #{e.message}", encryption_key)
+    end
+
+    def handle_file_chunk(sender_ip, message_data, active_transfers, callback, encryption_key)
+      chunk = JSON.parse(message_data)
+      transfer_id = chunk["id"]
+      transfer = active_transfers[transfer_id]
+      if transfer && transfer[:sender_ip] == sender_ip
+        chunk_data = Base64.strict_decode64(chunk["data"])
+        transfer[:temp_file].write(chunk_data)
+        transfer[:chunks_received] += 1
+        bytes_received = transfer[:temp_file].size
+        progress = (bytes_received.to_f / transfer[:file_size] * 100).round(2)
+        callback&.call(:progress, {
+                         transfer_id: transfer_id,
+                         sender_ip: sender_ip,
+                         file_name: transfer[:file_name],
+                         progress: progress,
+                         bytes_received: bytes_received,
+                         total_bytes: transfer[:file_size]
+                       })
+      end
+    rescue JSON::ParserError => e
+      send_error(sender_ip, "unknown", "Invalid chunk format: #{e.message}", encryption_key)
+    end
+
+    def handle_file_end(sender_ip, message_data, active_transfers, output_dir, encryption_key, callback)
+      end_data = JSON.parse(message_data)
+      transfer_id = end_data["id"]
+      transfer = active_transfers[transfer_id]
+      if transfer && transfer[:sender_ip] == sender_ip
+        transfer[:temp_file].close
+        calculated_checksum = calculate_file_checksum(transfer[:temp_file].path)
+        if calculated_checksum == transfer[:expected_checksum]
+          final_path = File.join(output_dir, transfer[:file_name])
+          FileUtils.mv(transfer[:temp_file].path, final_path)
+          ack_message = Lanet::Encryptor.prepare_message("#{FILE_ACK}#{transfer_id}", encryption_key)
+          @sender.send_to(sender_ip, ack_message)
+          callback&.call(:complete, {
+                           transfer_id: transfer_id,
+                           sender_ip: sender_ip,
+                           file_name: transfer[:file_name],
+                           file_path: final_path
+                         })
+        else
+          error_msg = "Checksum verification failed"
+          send_error(sender_ip, transfer_id, error_msg, encryption_key)
+          callback&.call(:error, {
+                           transfer_id: transfer_id,
+                           sender_ip: sender_ip,
+                           error: error_msg
+                         })
+        end
+        transfer[:temp_file].unlink
+        active_transfers.delete(transfer_id)
+      end
+    rescue JSON::ParserError => e
+      send_error(sender_ip, "unknown", "Invalid end marker format: #{e.message}", encryption_key)
+    end
+
+    def handle_file_error(sender_ip, message_data, active_transfers, callback)
+      error_data = JSON.parse(message_data)
+      transfer_id = error_data["id"]
+      if callback && active_transfers[transfer_id]
+        callback.call(:error, {
+                        transfer_id: transfer_id,
+                        sender_ip: sender_ip,
+                        error: error_data["message"]
+                      })
+        if active_transfers[transfer_id]
+          active_transfers[transfer_id][:temp_file].close
+          active_transfers[transfer_id][:temp_file].unlink
+          active_transfers.delete(transfer_id)
+        end
+      end
+    rescue JSON::ParserError
+      # Ignore malformed error messages
+    end
+
+    def cleanup_transfers(active_transfers)
+      active_transfers.each_value do |transfer|
+        transfer[:temp_file].close
+        begin
+          transfer[:temp_file].unlink
+        rescue StandardError
+          nil
+        end
+      end
+    end
+  end
+end

data/lib/lanet/scanner.rb

@@ -25,55 +25,54 @@ module Lanet
       8443 => "HTTPS-ALT"
     }.freeze
 
-    # Ports to check during scan
     QUICK_CHECK_PORTS = [80, 443, 22, 445, 139, 8080].freeze
 
     def initialize
       @hosts = []
       @mutex = Mutex.new
+      @arp_cache = {}
     end
 
-    # Scan network and return active hosts
     def scan(cidr, timeout = 1, max_threads = 32, verbose = false)
       @verbose = verbose
       @timeout = timeout
-
-      # Clear previous scan results
       @hosts = []
-
-      # Get the range of IP addresses to scan
       range = IPAddr.new(cidr).to_range
-
-      # Create a queue of IPs to scan
       queue = Queue.new
       range.each { |ip| queue << ip.to_s }
-
       total_ips = queue.size
       completed = 0
 
-      # Pre-populate ARP table to improve MAC address resolution
-      update_arp_table(cidr, max_threads)
+      # Initial ARP cache population
+      @arp_cache = parse_arp_table
 
-      # Create worker threads to process the queue
       threads = Array.new([max_threads, total_ips].min) do
         Thread.new do
-          while (ip = begin
-            queue.pop(true)
-          rescue ThreadError
-            nil
-          end)
+          loop do
+            begin
+              ip = queue.pop(true)
+            rescue ThreadError
+              break
+            end
             scan_host(ip)
             @mutex.synchronize do
               completed += 1
               if total_ips < 100 || (completed % 10).zero? || completed == total_ips
-                print_progress(completed,
-                               total_ips)
+                print_progress(completed, total_ips)
               end
             end
           end
         end
       end
 
+      # Periodically update ARP cache
+      arp_updater = Thread.new do
+        while threads.any?(&:alive?)
+          sleep 5
+          @mutex.synchronize { @arp_cache = parse_arp_table }
+        end
+      end
+
       begin
         threads.each(&:join)
         print_progress(total_ips, total_ips)
@@ -82,6 +81,8 @@ module Lanet
       rescue Interrupt
         puts "\nScan interrupted. Returning partial results..."
         @verbose ? @hosts : @hosts.map { |h| h[:ip] }
+      ensure
+        arp_updater.kill if arp_updater.alive?
       end
     end
 
@@ -92,47 +93,62 @@ module Lanet
       print "\rScanning network: #{percent}% complete (#{completed}/#{total})"
     end
 
-    def update_arp_table(cidr, max_threads = 10)
-      # Use fast ping method to update ARP table
-      range = IPAddr.new(cidr).to_range
-      queue = Queue.new
-      range.each { |ip| queue << ip.to_s }
+    def parse_arp_table
+      cmd = RbConfig::CONFIG["host_os"] =~ /mswin|mingw|cygwin/ ? "arp -a" : "arp -a"
+      output = `#{cmd}`
+      arp_cache = {}
 
-      total = queue.size
-      processed = 0
+      case RbConfig::CONFIG["host_os"]
+      when /darwin/
+        output.each_line do |line|
+          next unless line =~ /\((\d+\.\d+\.\d+\.\d+)\) at ([0-9a-f:]+) on/
 
-      threads = Array.new([max_threads, total].min) do
-        Thread.new do
-          while (ip = begin
-            queue.pop(true)
-          rescue ThreadError
-            nil
-          end)
-            # Use system ping to update ARP table
-            system("ping -c 1 -W 1 #{ip} > /dev/null 2>&1 &")
-            sleep 0.01 # Small delay to prevent overwhelming the system
-            processed += 1
-          end
+          ip = ::Regexp.last_match(1)
+          mac = ::Regexp.last_match(2).downcase
+          arp_cache[ip] = mac unless mac == "(incomplete)"
         end
-      end
+      when /linux/
+        output.each_line do |line|
+          next unless line =~ /^(\d+\.\d+\.\d+\.\d+)\s+\w+\s+([0-9a-f:]+)\s+/
+
+          ip = ::Regexp.last_match(1)
+          mac = ::Regexp.last_match(2).downcase
+          arp_cache[ip] = mac unless mac == "00:00:00:00:00:00"
+        end
+      when /mswin|mingw|cygwin/
+        output.each_line do |line|
+          next unless line =~ /^\s*(\d+\.\d+\.\d+\.\d+)\s+([0-9a-f-]+)\s+/
 
-      # Wait for ping operations to complete
-      threads.each(&:join)
-      sleep 1 # Give the system time to update ARP table
+          ip = ::Regexp.last_match(1)
+          mac = ::Regexp.last_match(2).gsub("-", ":").downcase
+          arp_cache[ip] = mac
+        end
+      end
+      arp_cache
     end
 
     def scan_host(ip)
-      # Skip special addresses to save time
-      return if ip.end_with?(".0") && !ip.end_with?(".0.0") # Skip network addresses except 0.0.0.0
+      # Handle broadcast addresses immediately
+      if ip.end_with?(".255") || ip == "255.255.255.255"
+        host_info = { ip: ip, mac: "ff:ff:ff:ff:ff:ff", response_time: 0, detection_method: "Broadcast" }
+        if @verbose
+          host_info[:hostname] = "Broadcast"
+          host_info[:ports] = {}
+        end
+        @mutex.synchronize { @hosts << host_info }
+        return
+      end
+
+      # Skip network addresses
+      return if ip.end_with?(".0") && !ip.end_with?(".0.0")
 
-      # Use multiple methods to detect if a host is alive
       is_active = false
       detection_method = nil
       response_time = nil
       start_time = Time.now
       open_ports = []
 
-      # Method 1: Try TCP port scan (most reliable)
+      # TCP port scan
       tcp_result = tcp_port_scan(ip, QUICK_CHECK_PORTS)
       if tcp_result[:active]
         is_active = true
@@ -140,116 +156,86 @@ module Lanet
         open_ports = tcp_result[:open_ports]
       end
 
-      # Method 2: Try ICMP ping if TCP didn't work
-      unless is_active
-        ping_result = ping_check(ip)
-        if ping_result
-          is_active = true
-          detection_method = "ICMP"
-        end
+      # ICMP ping
+      if !is_active && ping_check(ip)
+        is_active = true
+        detection_method = "ICMP"
       end
 
-      # Method 3: If host is a common network device (e.g., router), check with UDP
-      if !is_active && (ip.end_with?(".1") || ip.end_with?(".254") || ip.end_with?(".255"))
-        udp_result = udp_check(ip)
-        if udp_result
-          is_active = true
-          detection_method = "UDP"
-        end
+      # UDP check for common network devices
+      if !is_active && (ip.end_with?(".1") || ip.end_with?(".254")) && udp_check(ip)
+        is_active = true
+        detection_method = "UDP"
       end
 
-      # Method 4: ARP Check - if we have a MAC, the host is likely active
+      # ARP check
       unless is_active
-        mac_address = get_mac_address(ip)
-        if mac_address && mac_address != "(incomplete)"
+        mac = get_mac_address(ip)
+        if mac && mac != "(incomplete)"
           is_active = true
           detection_method = "ARP"
         end
       end
 
-      # For broadcast addresses, always consider them active
-      if ip.end_with?(".255") || ip == "255.255.255.255"
-        is_active = true
-        detection_method = "Broadcast"
-      end
-
-      # Calculate response time
       response_time = ((Time.now - start_time) * 1000).round(2) if is_active
-
       return unless is_active
 
-      # For active hosts, collect more information if in verbose mode
-      host_info = {
-        ip: ip,
-        mac: get_mac_address(ip),
-        response_time: response_time,
-        detection_method: detection_method
-      }
+      host_info = { ip: ip, mac: get_mac_address(ip), response_time: response_time, detection_method: detection_method }
 
       if @verbose
-        # For verbose mode, try to resolve hostname
-        begin
-          Timeout.timeout(1) do
-            host_info[:hostname] = Resolv.getname(ip)
-          end
-        rescue Resolv::ResolvError, Timeout::Error
-          host_info[:hostname] = "Unknown"
+        host_info[:hostname] = begin
+          Timeout.timeout(1) { Resolv.getname(ip) }
+        rescue StandardError
+          "Unknown"
         end
-
-        # For verbose mode, scan more ports if TCP detection method was successful
         if detection_method == "TCP"
           extra_ports = tcp_port_scan(ip, COMMON_PORTS.keys - QUICK_CHECK_PORTS)[:open_ports]
           open_ports += extra_ports
         end
-
         host_info[:ports] = open_ports.map { |port| [port, COMMON_PORTS[port] || "Unknown"] }.to_h
       end
 
       @mutex.synchronize { @hosts << host_info }
-    rescue StandardError => e
-      puts "\nError scanning host #{ip}: #{e.message}" if $DEBUG
     end
 
     def tcp_port_scan(ip, ports)
       open_ports = []
       is_active = false
+      threads = ports.map do |port|
+        Thread.new(port) do |p|
+          Timeout.timeout(@timeout) do
+            socket = TCPSocket.new(ip, p)
+            Thread.current[:open] = p
+            socket.close
+          end
+        rescue Errno::ECONNREFUSED
+          Thread.current[:active] = true
+        rescue StandardError
+          # Port closed or filtered
+        end
+      end
 
-      ports.each do |port|
-        Timeout.timeout(@timeout) do
-          socket = TCPSocket.new(ip, port)
+      threads.each do |thread|
+        thread.join
+        if thread[:open]
+          open_ports << thread[:open]
+          is_active = true
+        elsif thread[:active]
           is_active = true
-          open_ports << port
-          socket.close
         end
-      rescue Errno::ECONNREFUSED
-        # Connection refused means host is up but port is closed
-        is_active = true
-      rescue StandardError
-        # Other errors mean port is probably closed or filtered
       end
 
       { active: is_active, open_ports: open_ports }
     end
 
     def ping_check(ip)
-      cmd = case RbConfig::CONFIG["host_os"]
-            when /darwin/
-              "ping -c 1 -W 1 #{ip}"
-            when /linux/
-              "ping -c 1 -W 1 #{ip}"
-            when /mswin|mingw|cygwin/
-              "ping -n 1 -w 1000 #{ip}"
-            else
-              "ping -c 1 -W 1 #{ip}"
-            end
-
+      cmd = RbConfig::CONFIG["host_os"] =~ /mswin|mingw|cygwin/ ? "ping -n 1 -w 1000 #{ip}" : "ping -c 1 -W 1 #{ip}"
       system("#{cmd} > /dev/null 2>&1")
       $CHILD_STATUS.exitstatus.zero?
     end
 
     def udp_check(ip)
       common_udp_ports = [53, 67, 68, 123, 137, 138, 1900, 5353]
-
       common_udp_ports.each do |port|
         Timeout.timeout(0.5) do
           socket = UDPSocket.new
@@ -259,35 +245,15 @@ module Lanet
           return true
         end
       rescue Errno::ECONNREFUSED
-        return true # Connection refused means host is up
+        return true
       rescue StandardError
-        # Try next port
+        next
       end
       false
     end
 
     def get_mac_address(ip)
-      return "ff:ff:ff:ff:ff:ff" if ip.end_with?(".255") # Special case for broadcast
-
-      # Get MAC from ARP table
-      cmd = case RbConfig::CONFIG["host_os"]
-            when /darwin/
-              "arp -n #{ip}"
-            when /linux/
-              "arp -n #{ip}"
-            when /mswin|mingw|cygwin/
-              "arp -a #{ip}"
-            else
-              "arp -n #{ip}"
-            end
-
-      output = `#{cmd}`
-
-      if output =~ /([0-9a-fA-F]{1,2}[:-][0-9a-fA-F]{1,2}[:-][0-9a-fA-F]{1,2}[:-][0-9a-fA-F]{1,2}[:-][0-9a-fA-F]{1,2}[:-][0-9a-fA-F]{1,2})/
-        ::Regexp.last_match(1).downcase
-      else
-        "(incomplete)"
-      end
+      @mutex.synchronize { @arp_cache[ip] || "(incomplete)" }
     end
   end
 end

data/lib/lanet/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Lanet
-  VERSION = "0.2.0"
+  VERSION = "0.3.0"
 end