From b5dcc0eb1dbd1c95f60a588ce3d8af7435e3e504 Mon Sep 17 00:00:00 2001
From: sinn3r <wei_chen@rapid7.com>
Date: Tue, 18 Feb 2014 12:43:11 -0600
Subject: [PATCH] Make several changes.

Some important changes:

* Uses optparse to parse argumnets
* Prevent file handle leaks
---
 tools/makeiplist.rb | 112 ++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 102 insertions(+), 10 deletions(-)
 mode change 100644 => 100755 tools/makeiplist.rb

diff --git a/tools/makeiplist.rb b/tools/makeiplist.rb
old mode 100644
new mode 100755
index 5c16328e1b1c..f9bfe7d641f9
--- a/tools/makeiplist.rb
+++ b/tools/makeiplist.rb
@@ -1,9 +1,18 @@
 #!/usr/bin/env ruby
+
 #
+# This script takes a list of ranges and converts it to a per line ip list.
+# Demonstration:
+# echo 192.168.100.0-50 >> rangelist.txt
+# echo 192.155-156.0.1 >> rangelist.txt
+# echo 192.168.200.0/25 >> rangelist.txt
+# ruby tools/makeiplist.rb
 #
-# This script takes a list of ranges and converts it to a per line ip list
+# Author:
+# mubix
 #
 
+
 msfbase = __FILE__
 while File.symlink?(msfbase)
   msfbase = File.expand_path(File.readlink(msfbase), File.dirname(msfbase))
@@ -12,16 +21,99 @@
 $:.unshift(File.expand_path(File.join(File.dirname(msfbase), '..', 'lib')))
 require 'msfenv'
 require 'rex'
+require 'optparse'
+
+
+class OptsConsole
+  def self.parse(args)
+    options = {'output' => 'iplist.txt'}
+
+    opts = OptionParser.new do |opts|
+      opts.banner = %Q|This script takes a list of ranges and converts it to a per line ip list.
+Usage: #{__FILE__} [options]|
+
+      opts.separator ""
+      opts.separator "Specific options:"
+
+      opts.on("-i", '-i <filename>', "Input file") do |v|
+        options['input'] = v.to_s
+      end
+
+      opts.on("-o", '-o <filename>', "(Optional) Output file. Default: iplist.txt") do |v|
+        options['output'] = v.to_s
+      end
 
+      opts.separator ""
+      opts.separator "Common options:"
 
-f = File.open('rangelist.txt', 'r')
-w = File.open('iplist.txt', 'a')
+      opts.on_tail("-h", "--help", "Show this message") do
+        puts opts
+        exit
+      end
+    end
 
-f.each_line do |range|
-        ips = Rex::Socket::RangeWalker.new(range)
-        ips.each do |ip|
-                w.puts ip
-        end
+    begin
+      opts.parse!(args)
+
+      unless ::File.exists?(options['input'])
+        raise OptionParser::InvalidArgument, "Not found: #{options['input']}"
+      end
+    rescue OptionParser::InvalidOption
+      puts "[*] Invalid option, try -h for usage"
+      exit
+    rescue OptionParser::InvalidArgument => e
+      puts "[*] #{e.message}"
+      exit
+    end
+
+    if options.empty?
+      puts "[*] No options specified, try -h for usage"
+      exit
+    end
+
+    options
+  end
 end
-f.close
-w.close
+
+
+#
+# Prints IPs
+#
+def make_list(in_f, out_f)
+  in_f.each_line do |range|
+    ips = Rex::Socket::RangeWalker.new(range)
+    ips.each do |ip|
+      out_f.puts ip
+      end
+    end
+end
+
+
+#
+# Returns file handles
+#
+def load_files(in_f, out_f)
+  handle_in = ::File.open(in_f, 'r')
+
+  # Output file not found, assuming we should create one automatically
+  ::File.open(out_f, 'w') {} unless ::File.exists?(out_f)
+
+  handle_out = ::File.open(out_f, 'a')
+
+  return handle_in, handle_out
+end
+
+
+options = OptsConsole.parse(ARGV)
+in_f, out_f = load_files(options['input'], options['output'])
+
+begin
+  puts "[*] Generating list at #{options['output']}"
+  make_list(in_f, out_f)
+ensure
+  # Always makes sure the file descriptors are closed
+  in_f.close
+  out_f.close
+end
+
+puts "[*] Done."