Add UDP support

lib/iptables.rb

@@ -1,4 +1,5 @@
 require 'iptables/protocol/tcp'
+require 'iptables/protocol/udp'
 require 'iptables/rule'
 
 module IPTables

lib/iptables/protocol/udp.rb

@@ -0,0 +1,31 @@
+module IPTables
+  module Protocol
+    module UDP
+      attr_reader :source_port, :destination_port
+
+      def self.extended(base)
+        base.mod_opts[:protocol] = []
+      end
+
+      def source_port=(value)
+        value = value.to_i
+        unless value > 0
+          raise "UDP source port must be a valid integer greater than 0"
+        end
+
+        @source_port = value
+        @mod_opts[:protocol] << "--sport" << source_port
+      end
+
+      def destination_port=(value)
+        value = value.to_i
+        unless value > 0
+          raise "UDP destination port must be a valid integer greater than 0"
+        end
+
+        @destination_port = value
+        @mod_opts[:protocol] << "--dport" << destination_port
+      end
+    end
+  end
+end

lib/iptables/rule.rb

@@ -11,6 +11,7 @@ def initialize
     def protocol=(value)
       protocols = {
         :tcp => IPTables::Protocol::TCP,
+        :udp => IPTables::Protocol::UDP,
       }
 
       value = value.to_sym unless value.is_a? Symbol

spec/iptables/protocol_tcp_spec.rb

@@ -18,10 +18,6 @@
     it { should respond_to(:tcp_flags=).with(1).argument }
     it { should respond_to(:syn).with(0).arguments }
     it { should respond_to(:syn=).with(1).arguments }
-    it { should respond_to(:tcp_option).with(0).arguments }
-    it { should respond_to(:tcp_option=).with(1).argument }
-    it { should respond_to(:mss).with(0).arguments }
-    it { should respond_to(:mss=).with(1).argument }
 
     context "when creating a rule with source_port" do
       subject do

spec/iptables/protocol_udp_spec.rb

@@ -0,0 +1,50 @@
+require 'spec_helper'
+
+describe IPTables::Rule do
+  describe "when protocol is udp" do
+    subject do
+      rule = IPTables::Rule.new
+      rule.protocol = :udp
+      rule
+    end
+
+    its(:protocol) { should == :udp }
+
+    it { should respond_to(:source_port).with(0).arguments }
+    it { should respond_to(:source_port=).with(1).argument }
+    it { should respond_to(:destination_port).with(0).arguments }
+    it { should respond_to(:destination_port=).with(1).argument }
+
+    context "and creating a rule with source_port" do
+      subject do
+        rule = IPTables::Rule.new
+        rule.chain = :input
+        rule.target = :accept
+        rule.protocol = :udp
+        rule.source_port = 80
+        rule
+      end
+
+      its(:source_port) { should == 80 }
+      its(:to_iptables) {
+        should == "-A INPUT -p udp --sport 80 -j ACCEPT"
+      }
+    end
+
+    context "and creating a rule with destination_port" do
+      subject do
+        rule = IPTables::Rule.new
+        rule.chain = :input
+        rule.target = :accept
+        rule.protocol = :udp
+        rule.destination_port = 443
+        rule
+      end
+
+      its(:destination_port) { should == 443 }
+      its(:to_iptables) {
+        should == "-A INPUT -p udp --dport 443 -j ACCEPT"
+      }
+    end
+  end
+end