zsteg 0.0.0 → 0.0.1

data/lib/zsteg/cli.rb

@@ -1,11 +1,8 @@
 require 'optparse'
-require 'awesome_print'
 
 module ZSteg
   class CLI
     DEFAULT_ACTIONS = %w'check'
-    DEFAULT_LIMIT   = 256
-    DEFAULT_ORDER   = 'auto'
 
     def initialize argv = ARGV
       @argv = argv
@@ -15,30 +12,35 @@ module ZSteg
       @actions = []
       @options = {
         :verbose => 0,
-        :limit => DEFAULT_LIMIT,
-        :bits  => [1,2,3,4],
-        :order => DEFAULT_ORDER
+        :limit => Checker::DEFAULT_LIMIT,
+        :order => Checker::DEFAULT_ORDER
       }
       optparser = OptionParser.new do |opts|
-        opts.banner = "Usage: zsteg [options] filename.png"
+        opts.banner = "Usage: zsteg [options] filename.png [param_string]"
         opts.separator ""
 
-        opts.on("-c", "--channels X", /[rgba,]+/,
+        opts.on("-c", "--channels X", /[rgba,1-8]+/,
                 "channels (R/G/B/A) or any combination, comma separated",
-                "valid values: r,g,b,a,rg,rgb,bgr,rgba,..."
+                "valid values: r,g,b,a,rg,bgr,rgba,r3g2b3,..."
         ){ |x| @options[:channels] = x.split(',') }
 
         opts.on("-l", "--limit N", Integer,
-                "limit bytes checked, 0 = no limit (default: #{DEFAULT_LIMIT})"
+                "limit bytes checked, 0 = no limit (default: #{@options[:limit]})"
         ){ |n| @options[:limit] = n }
 
-        opts.on("-b", "--bits N", /[\d,-]+/,
-                "number of bits (1..8), single value or '1,3,5' or '1-8'") do |n|
-          if n['-']
-            @options[:bits] = Range.new(*n.split('-').map(&:to_i)).to_a
-          else
-            @options[:bits] = n.split(',').map(&:to_i)
+        opts.on("-b", "--bits N", "number of bits, single int value or '1,3,5' or range '1-8'",
+                "advanced: specify individual bits like '00001110' or '0x88'"
+        ) do |x|
+          a = []
+          x.split(',').each do |x1|
+            if x1['-']
+              t = x1.split('-')
+              a << Range.new(parse_bits(t[0]), parse_bits(t[1])).to_a
+            else
+              a << parse_bits(x1)
+            end
           end
+          @options[:bits] = a.flatten.uniq
         end
 
         opts.on "--lsb", "least significant BIT comes first" do
@@ -48,8 +50,20 @@ module ZSteg
           @options[:bit_order] = :msb
         end
 
+        opts.on "-P", "--prime", "analyze/extract only prime bytes/pixels" do
+          @options[:prime] = true
+        end
+#        opts.on "--pixel-align", "pixel-align hidden data (EasyBMP)" do
+#          @options[:pixel_align] = true
+#        end
+
+        opts.on "-a", "--all", "try all known methods" do
+          @options[:prime] = :all
+          @options[:order] = :all
+        end
+
         opts.on("-o", "--order X", /all|auto|[bxy,]+/i,
-                "pixel iteration order (default: '#{DEFAULT_ORDER}')",
+                "pixel iteration order (default: '#{@options[:order]}')",
                 "valid values: ALL,xy,yx,XY,YX,xY,Xy,bY,...",
         ){ |x| @options[:order] = x.split(',') }
 
@@ -64,6 +78,11 @@ module ZSteg
         opts.on "-q", "--quiet", "Silent any warnings (can be used multiple times)" do |v|
           @options[:verbose] -= 1
         end
+        opts.on "-C", "--[no-]color", "Force (or disable) color output (default: auto)" do |x|
+          Sickill::Rainbow.enabled = x
+        end
+        opts.separator "\nPARAMS SHORTCUT\n"+
+          "\tzsteg fname.png 2b,b,lsb,xy  ==>  --bits 2 --channel b --lsb --order xy"
       end
 
       if (argv = optparser.parse(@argv)).empty?
@@ -73,12 +92,19 @@ module ZSteg
 
       @actions = DEFAULT_ACTIONS if @actions.empty?
 
+      argv.each do |arg|
+        if arg[','] && !File.exist?(arg)
+          @options.merge!(decode_param_string(arg))
+          argv.delete arg
+        end
+      end
+
       argv.each_with_index do |fname,idx|
         if argv.size > 1 && @options[:verbose] >= 0
           puts if idx > 0
           puts "[.] #{fname}".green
         end
-        @fname = fname
+        next unless @img=load_image(@fname=fname)
 
         @actions.each do |action|
           if action.is_a?(Array)
@@ -93,39 +119,70 @@ module ZSteg
       # prevents a 'Broken pipe - <STDOUT> (Errno::EPIPE)' message
     end
 
-    ###########################################################################
-    # actions
-
-    def check
-      Checker.new(@fname, @options).check
+    def load_image fname
+      if File.directory?(fname)
+        puts "[?] #{fname} is a directory".yellow
+      else
+        ZPNG::Image.load(fname)
+      end
+    rescue ZPNG::Exception, Errno::ENOENT
+      puts "[!] #{$!.inspect}".red
     end
 
-    def extract name
-      if ['extradata', 'data after IEND'].include?(name)
-        img = ZPNG::Image.load(@fname)
-        print img.extradata
-        return
+    def parse_bits x
+      case x
+        when '1', 1             # catch NOT A BINARY MASK early
+          1
+        when /^0x[0-9a-f]+$/i   # hex,     mask
+          0x100 + x.to_i(16)
+        when /^(?:0b)?[01]+$/i  # binary,  mask
+          0x100 + x.to_i(2)
+        when /^\d+$/            # decimal, number of bits
+          x.to_i
+        else
+          raise "invalid bits value: #{x.inspect}"
       end
+    end
 
+    def decode_param_string s
       h = {}
-      name.split(',').each do |x|
+      s.split(',').each do |x|
         case x
         when 'lsb'
           h[:bit_order] = :lsb
         when 'msb'
           h[:bit_order] = :msb
-        when /(\d)b/
-          h[:bits] = $1.to_i
+        when /^(\d)b$/, /^b(\d)$/
+          h[:bits] = parse_bits($1)
         when /\A[rgba]+\Z/
-          h[:channels] = x.split('')
-        when /\Axy|yx\Z/i
+          h[:channels] = [x]
+        when /\Axy|yx|yb|by\Z/i
           h[:order] = x
+        when 'prime'
+          h[:prime] = true
         else
           raise "uknown param #{x.inspect}"
         end
       end
-      h[:limit] = @options[:limit] if @options[:limit] != DEFAULT_LIMIT
-      print Extractor.new(@fname, @options).extract(h)
+      h
+    end
+
+    ###########################################################################
+    # actions
+
+    def check
+      Checker.new(@img, @options).check
+    end
+
+    def extract name
+      if ['extradata', 'data after IEND'].include?(name)
+        print @img.extradata
+        return
+      end
+
+      h = decode_param_string name
+      h[:limit] = @options[:limit] if @options[:limit] != Checker::DEFAULT_LIMIT
+      print Extractor.new(@img, @options).extract(h)
     end
 
   end

data/lib/zsteg/extractor/byte_extractor.rb

@@ -5,22 +5,23 @@ module ZSteg
     module ByteExtractor
 
       def byte_extract params = {}
-        limit = params[:limit].to_i
-        limit = 2**32 if limit <= 0
-
-        bits = params[:bits]
-        raise "invalid bits value #{bits.inspect}" unless (1..8).include?(bits)
-        mask = 2**bits - 1
+        masks = bits2masks params[:bits]
 
+        if params[:prime]
+          pregenerate_primes(
+            :max   => @image.scanlines[0].size * @image.height,
+            :count => (@limit*8.0/masks.size).ceil
+          )
+        end
 
         data = ''.force_encoding('binary')
         a = []
-        byte_iterator(params[:order]) do |x,y|
+        byte_iterator(params) do |x,y|
           sl = @image.scanlines[y]
 
-          value = sl.decoded_bytes[x].ord
-          bits.times do |bidx|
-            a << ((value & (1<<(bits-bidx-1))) == 0 ? 0 : 1)
+          value = sl.decoded_bytes.getbyte(x)
+          masks.each do |mask|
+            a << ((value & mask) == 0 ? 0 : 1)
           end
 
           if a.size >= 8
@@ -32,8 +33,8 @@ module ZSteg
             end
             #printf "[d] %02x %08b\n", byte, byte
             data << byte.chr
-            if data.size >= limit
-              print "[limit #{params[:limit]}]".gray if @verbose > 1
+            if data.size >= @limit
+              print "[limit #@limit]".gray if @verbose > 1
               break
             end
           end
@@ -51,7 +52,8 @@ module ZSteg
       # ...
       # 'xY': b=0,  y=MAX; b=1,    y=MAX; b=2,    y=MAX; ...
       # 'XY': b=MAX,y=MAX; b=MAX-1,y=MAX; b=MAX-2,y=MAX; ...
-      def byte_iterator type = nil
+      def byte_iterator params
+        type = params[:order]
         if type.nil? || type == 'auto'
           type = @image.format == :bmp ? 'bY' : 'by'
         end
@@ -59,6 +61,7 @@ module ZSteg
 
         sl0 = @image.scanlines.first
 
+        # XXX don't try to run it on interlaced PNGs!
         x0,x1,xstep =
           if type.index('b')
             [0, sl0.decoded_bytes.size-1, 1]
@@ -73,19 +76,31 @@ module ZSteg
             [@image.height-1, 0, -1]
           end
 
+        # cannot join these lines from ByteExtractor and ColorExtractor into
+        # one method for performance reason:
+        #   it will require additional yield() for EACH BYTE iterated
+
         if type[0,1].downcase == 'b'
           # ROW iterator
-          y0.step(y1,ystep) do |y|
-            x0.step(x1,xstep) do |x|
-              yield x,y
-            end
+          if params[:prime]
+            idx = 0
+            y0.step(y1,ystep){ |y| x0.step(x1,xstep){ |x|
+              yield(x,y) if @primes.include?(idx)
+              idx += 1
+            }}
+          else
+            y0.step(y1,ystep){ |y| x0.step(x1,xstep){ |x| yield(x,y) }}
           end
         else
           # COLUMN iterator
-          x0.step(x1,xstep) do |x|
-            y0.step(y1,ystep) do |y|
-              yield x,y
-            end
+          if params[:prime]
+            idx = 0
+            x0.step(x1,xstep){ |x| y0.step(y1,ystep){ |y|
+              yield(x,y) if @primes.include?(idx)
+              idx += 1
+            }}
+          else
+            x0.step(x1,xstep){ |x| y0.step(y1,ystep){ |y| yield(x,y) }}
           end
         end
       end

data/lib/zsteg/extractor/color_extractor.rb

@@ -4,40 +4,61 @@ module ZSteg
     module ColorExtractor
 
       def color_extract params = {}
-        channels = (Array(params[:channels]) + Array(params[:channel])).compact
+        channels = Array(params[:channels])
+        #pixel_align = params[:pixel_align]
 
-        limit = params[:limit].to_i
-        limit = 2**32 if limit <= 0
+        ch_masks = []
+        case channels.first.size
+        when 1
+          # ['r', 'g', 'b']
+          channels.each{ |c| ch_masks << [c[0], bits2masks(params[:bits])] }
+        when 2
+          # ['r3', 'g2', 'b3']
+          channels.each{ |c| ch_masks << [c[0], bits2masks(c[1].to_i)] }
+        else
+          raise "invalid channels: #{channels.inspect}"
+        end
 
-        bits = params[:bits]
-        raise "invalid bits value #{bits.inspect}" unless (1..8).include?(bits)
-        mask = 2**bits - 1
+        # total number of bits = sum of all channels bits
+        nbits = ch_masks.map{ |x| x[1].size }.inject(&:+)
 
+        if params[:prime]
+          pregenerate_primes(
+            :max   => @image.width * @image.height,
+            :count => (@limit*8.0/nbits/channels.size).ceil
+          )
+        end
 
         data = ''.force_encoding('binary')
         a = []
-        coord_iterator(params[:order]) do |x,y|
-          color = @image[x,y]
+        #puts
+        catch :limit do
+          coord_iterator(params) do |x,y|
+            color = @image[x,y]
 
-          channels.each do |c|
-            value = color.send(c)
-            bits.times do |bidx|
-              a << ((value & (1<<(bits-bidx-1))) == 0 ? 0 : 1)
+            ch_masks.each do |c,masks|
+              value = color.send(c)
+              masks.each do |mask|
+                a << ((value & mask) == 0 ? 0 : 1)
+              end
             end
-          end
+            #p [x,y,a.size,a]
 
-          if a.size >= 8
-            byte = 0
-            if params[:bit_order] == :msb
-              8.times{ |i| byte |= (a.shift<<i)}
-            else
-              8.times{ |i| byte |= (a.shift<<(7-i))}
-            end
-            #printf "[d] %02x %08b\n", byte, byte
-            data << byte.chr
-            if data.size >= limit
-              print "[limit #{params[:limit]}]".gray if @verbose > 1
-              break
+            while a.size >= 8
+              byte = 0
+              #puts a.join
+              if params[:bit_order] == :msb
+                8.times{ |i| byte |= (a.shift<<i)}
+              else
+                8.times{ |i| byte |= (a.shift<<(7-i))}
+              end
+              #printf "[d] %02x %08b\n", byte, byte
+              data << byte.chr
+              if data.size >= @limit
+                print "[limit #@limit]".gray if @verbose > 1
+                throw :limit
+              end
+              #a.clear if pixel_align
             end
           end
         end
@@ -54,7 +75,8 @@ module ZSteg
       # ...
       # 'xY': x=0,  y=MAX; x=1,    y=MAX; x=2,    y=MAX; ...
       # 'XY': x=MAX,y=MAX; x=MAX-1,y=MAX; x=MAX-2,y=MAX; ...
-      def coord_iterator type = nil
+      def coord_iterator params
+        type = params[:order]
         if type.nil? || type == 'auto'
           type = @image.format == :bmp ? 'xY' : 'xy'
         end
@@ -74,19 +96,31 @@ module ZSteg
             [@image.height-1, 0, -1]
           end
 
+        # cannot join these lines from ByteExtractor and ColorExtractor into
+        # one method for performance reason:
+        #   it will require additional yield() for EACH BYTE iterated
+
         if type[0,1].downcase == 'x'
           # ROW iterator
-          y0.step(y1,ystep) do |y|
-            x0.step(x1,xstep) do |x|
-              yield x,y
-            end
+          if params[:prime]
+            idx = 0
+            y0.step(y1,ystep){ |y| x0.step(x1,xstep){ |x|
+              yield(x,y) if @primes.include?(idx)
+              idx += 1
+            }}
+          else
+            y0.step(y1,ystep){ |y| x0.step(x1,xstep){ |x| yield(x,y) }}
           end
         else
           # COLUMN iterator
-          x0.step(x1,xstep) do |x|
-            y0.step(y1,ystep) do |y|
-              yield x,y
-            end
+          if params[:prime]
+            idx = 0
+            x0.step(x1,xstep){ |x| y0.step(y1,ystep){ |y|
+              yield(x,y) if @primes.include?(idx)
+              idx += 1
+            }}
+          else
+            x0.step(x1,xstep){ |x| y0.step(y1,ystep){ |y| yield(x,y) }}
           end
         end
       end

data/lib/zsteg/extractor.rb

@@ -1,3 +1,6 @@
+require 'prime'
+require 'set'
+
 module ZSteg
   class Extractor
 
@@ -7,15 +10,47 @@ module ZSteg
     # image can be either filename or ZPNG::Image
     def initialize image, params = {}
       @image = image.is_a?(ZPNG::Image) ? image : ZPNG::Image.load(image)
-      @verbose = params[:verbose]
+      @verbose = params[:verbose] || 0
     end
 
     def extract params = {}
+      @limit = params[:limit].to_i
+      @limit = 2**32 if @limit <= 0
+
       if params[:order] =~ /b/i
         byte_extract params
       else
         color_extract params
       end
     end
+
+    def pregenerate_primes h
+      @primes ||= Set.new
+      return if @primes.size >= h[:count]
+
+      count = h[:count]
+      Prime.each(h[:max]) do |prime|
+        @primes << prime
+        break if @primes.size >= count
+      end
+    end
+
+    def bits2masks bits
+      masks = []
+      if (1..8).include?(bits)
+        # number of bits
+        bits.times do |i|
+          masks << (1<<(bits-i-1))
+        end
+      else
+        # mask
+        bits &= 0xff
+        8.times do |i|
+          mask = (1<<(bits-i-1))
+          masks << mask if (bits & mask) != 0
+        end
+      end
+      masks
+    end
   end
 end