lib/packetgen/packet.rb in packetgen-1.0.1 vs lib/packetgen/packet.rb in packetgen-1.1.0

@@ -1,5 +1,6 @@
+# coding: utf-8
 # 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.
 require 'pcaprub'
@@ -161,29 +162,11 @@
     # @raise [ArgumentError] unknown protocol
     def add(protocol, options={})
       klass = check_protocol(protocol)
 
       header = klass.new(options)
-      prev_header = @headers.last
-      if prev_header
-        binding = prev_header.class.known_headers[klass]
-        if binding.nil?
-          msg = "#{prev_header.class} knowns no layer association with #{protocol}. "
-          msg << "Try #{prev_header.class}.bind_layer(PacketGen::Header::#{protocol}, "
-          msg << "#{prev_header.protocol_name.downcase}_proto_field: "
-          msg << "value_for_#{protocol.downcase})"
-          raise ArgumentError, msg
-        end
-        prev_header[binding.key].read binding.value
-        prev_header.body = header
-      end
-      header.packet = self
-      @headers << header
-      unless respond_to? protocol.downcase
-        self.class.class_eval "def #{protocol.downcase}(arg=nil);" \
-                              "header('#{protocol}', arg); end"
-      end
+      add_header header
       self
     end
 
     # Check if a protocol header is embedded in packet
     # @return [Boolean]
@@ -255,10 +238,38 @@
         type = @headers.first.protocol_name
         raise WireError, "don't known how to send a #{type} packet on wire"
       end
     end
 
+    # Encapulate another packet in +self+
+    # @param [Packet] other
+    # @return [self] +self+ with new headers from +other+
+    # @since 1.1.0
+    def encapsulate(other)
+      other.headers.each { |h| add_header h }
+    end
+
+    # Remove headers from +self+
+    # @param [Array<Header>] headers
+    # @return [self] +self+ with some headers removed
+    # @raise [FormatError] any headers not in +self+
+    # @raise [FormatError] removed headers result in an unknown binding
+    # @since 1.1.0
+    def decapsulate(*headers)
+      headers.each do |header|
+        idx = @headers.index(header)
+        raise FormatError, 'header not in packet!' if idx.nil?
+
+        prev_header = idx > 0 ? @headers[idx - 1] : nil
+        next_header = (idx+1) < @headers.size ? @headers[idx + 1] : nil
+        @headers.delete_at(idx)
+        add_header(next_header, prev_header) if prev_header and next_header
+      end
+    rescue ArgumentError => ex
+      raise FormatError, ex.message
+    end
+
     # @return [String]
     def inspect
       str = Inspect.dashed_line(self.class)
       @headers.each do |header|
         str << header.inspect
@@ -272,29 +283,27 @@
       to_s == other.to_s
     end
 
     private
 
-    # @overload header(protocol, layer=1)
-    #  @param [String] protocol
+    # @overload header(klass, layer=1)
+    #  @param [Class] klass
     #  @param [Integer] layer
-    # @overload header(protocol, options)
-    #  @param [String] protocol
+    # @overload header(klass, options={})
+    #  @param [String] klass
     #  @param [Hash] options
+    #  @raise [ArgumentError] unknown option
     # @return [Header::Base]
-    # @raise [ArgumentError] unknown protocol
-    def header(protocol, arg)
-      klass = check_protocol protocol
-
+    def header(klass, arg)
       headers = @headers.select { |h| h.is_a? klass }
       layer = arg.is_a?(Integer) ? arg : 1
       header = headers[layer - 1]
 
       if arg.is_a? Hash
         arg.each do |key, value|
           unless header.respond_to? "#{key}="
-            raise ArgumentError, "unknown #{key} attribute for #{header.class}"
+            raise ArgumentError, "unknown #{key} attribute for #{klass}"
           end
           header.send "#{key}=", value
         end
       end
 
@@ -303,15 +312,39 @@
 
     # check if protocol is known
     # @param [String] protocol
     # @raise [ArgumentError] unknown protocol
     def check_protocol(protocol)
-      unless Header.const_defined? protocol
-        raise ArgumentError, "unknown #{protocol} protocol"
-      end
-      klass = Header.const_get(protocol)
-      raise ArgumentError, "unknown #{protocol} protocol" unless klass.is_a? Class
+      klass = Header.get_header_class_by_name(protocol)
+      raise ArgumentError, "unknown #{protocol} protocol" if klass.nil?
       klass
+    end
+
+    # Add a header to packet
+    # @param [Header::HeaderMethods] header
+    # @param [Header::HeaderMethods] previous_header
+    # @return [void]
+    def add_header(header, previous_header=nil)
+      protocol = header.protocol_name
+      prev_header = previous_header || @headers.last
+      if prev_header
+        binding = prev_header.class.known_headers[header.class]
+        if binding.nil?
+          msg = "#{prev_header.class} knowns no layer association with #{protocol}. "
+          msg << "Try #{prev_header.class}.bind_layer(PacketGen::Header::#{protocol}, "
+          msg << "#{prev_header.protocol_name.downcase}_proto_field: "
+          msg << "value_for_#{protocol.downcase})"
+          raise ArgumentError, msg
+        end
+        prev_header[binding.key].read binding.value
+        prev_header.body = header
+      end
+      header.packet = self
+      @headers << header unless previous_header
+      unless respond_to? protocol.downcase
+        self.class.class_eval "def #{protocol.downcase}(arg=nil);" \
+                              "header(#{header.class}, arg); end"
+      end
     end
   end
 end
 
 require_relative 'header'