lib/packetgen/pcapng/shb.rb in packetgen-3.1.3 vs lib/packetgen/pcapng/shb.rb in packetgen-3.1.4

@@ -1,12 +1,12 @@
+# frozen_string_literal: true
+
 # This file is part of PacketGen
 # See https://github.com/sdaubert/packetgen for more informations
 # Copyright (C) 2016 Sylvain Daubert <[e-mail]>
 # This program is published under MIT license.
 
-# frozen_string_literal: true
-
 module PacketGen
   module PcapNG
     # {SHB} represents a Section Header Block (SHB) of a pcapng file.
     #
     # == SHB Definition
@@ -87,43 +87,24 @@
 
       # Reads a String or a IO to populate the object
       # @param [::String,IO] str_or_io
       # @return [self]
       def read(str_or_io)
-        io = if str_or_io.respond_to? :read
-               str_or_io
-             else
-               StringIO.new(force_binary(str_or_io.to_s))
-             end
+        io = to_io(str_or_io)
         return self if io.eof?
 
-        type_str = io.read(4)
-        unless type_str == PcapNG::SHB_TYPE.to_s
-          type = type_str.unpack('H*').join
-          raise InvalidFileError, "Incorrect type (#{type})for Section Header Block"
-        end
-
+        type_str = check_shb(io)
         block_len_str = io.read(4)
-
         magic_str = io.read(4)
-        case @endian
-        when :little
-          case magic_str
-          when MAGIC_LITTLE
-          when MAGIC_BIG
-            force_endianness :big
-          else
-            raise InvalidFileError, 'Incorrect magic for Section Header Block'
-          end
-        when :big
-          case magic_str
-          when MAGIC_BIG
-          when MAGIC_LITTLE
-            force_endianness :little
-          else
-            raise InvalidFileError, 'Incorrect magic for Section Header Block'
-          end
+
+        case magic_str
+        when MAGIC_LITTLE
+          force_endianness :little if @endian == :big
+        when MAGIC_BIG
+          force_endianness :big if @endian == :little
+        else
+          raise InvalidFileError, 'Incorrect magic for Section Header Block'
         end
 
         self[:type].read type_str
         self[:block_len].read block_len_str
         self[:magic].read magic_str
@@ -147,13 +128,12 @@
 
       # Return the object as a String
       # @return [String]
       def to_s
         body = @interfaces.map(&:to_s).join
-        unless self.section_len == SECTION_LEN_UNDEFINED
-          self.section_len = body.size
-        end
+        self.section_len = body.size unless self.section_len == SECTION_LEN_UNDEFINED
+
         pad_field :options
         recalc_block_len
         super + body
       end
 
@@ -166,9 +146,20 @@
         self[:magic] = Types::Int32.new(self[:magic].to_i, endian)
         self[:ver_major] = Types::Int16.new(self[:ver_major].to_i, endian)
         self[:ver_minor] = Types::Int16.new(self[:ver_minor].to_i, endian)
         self[:section_len] = Types::Int64.new(self[:section_len].to_i, endian)
         self[:block_len2] = Types::Int32.new(self[:block_len2].to_i, endian)
+      end
+
+      # Check io contains a SHB
+      # @param [IO] io
+      # @return [String] type string
+      def check_shb(io)
+        type_str = io.read(4)
+        return type_str if type_str == PcapNG::SHB_TYPE.to_s
+
+        type = type_str.unpack('H*').join
+        raise InvalidFileError, "Incorrect type (#{type})for Section Header Block"
       end
     end
   end
 end