Ruby script that reads/writes jpeg comment segments

jpgcom.rb

#!/usr/bin/env ruby
# To the extent possible under law, the author has waived all copyright
# and related or neighboring rights to this software by associating the
# CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/) with it.
#
# This script reads/writes COM (comment) segments in a jpeg stream.
# usage: ruby jpgcom.rb [-r] [-0|-1|...|-e|-f] [filename ...]
#
# When invoked with no argument, it reads a jpeg stream and writes the
# contents of its COM segments to stdout.  If there are multiple COM
# segments, they are simply concatinated.
# When -r option is specified, it reads the jpeg stream from the given
# file instead of stdin.
#
# Otherwise, it reads a jpeg stream and COM data from the given files,
# inserts the COM data to the jpeg stream, and writes the result to
# stdout. At least one file must be given. The first file must be the
# jpeg file and others are COM data to be inserted.  '-' means stdin.
# If the data does not fit within a segment limit (65,533 bytes),
# sequential multiple COM segments will be created.
#
# When one of -0, -1, ..., -f options is specified, it deals with APP0,
# APP1, ..., APP15 instead of COM.

class JPGCOM

  class Error < StandardError; end

  private

  def die(s)
    raise Error, s
  end

  def parse_jpeg(src)
    ret, pos = [], 0
    begin
      a = (src[pos,4] || '').unpack('nn')
      case a[0]
      when 0xffff then pos += 1; next
      when 0xffd8, 0xffd9 then a[1] = 0 # start of image, end of image
      when 0..0xff00, nil then die 'not in jpeg format'
      else die 'segment without length found' unless a[1] and a[1] >= 2
      end
      ret.push({:pos => pos, :mark => a[0], :len => a[1]})
      pos += 2 + a[1]
      case a[0] when 0xffda then # start of scan
        m = /(?=\xff[^\x00\xd0-\xd7])/mn.match(src, pos)
        pos = m ? m.begin(0) : die('unterminated entropy-coded segment')
      end
    end while pos < src.length
    ret
  end

  public

  def read(jpeg, target = 0xfffe)
    parse_jpeg(jpeg).each do |mark:0,pos:0,len:0|
      yield(jpeg[pos+4,len-2]) if mark == target
    end
  end

  def write(jpeg, data, target = 0xfffe)
    parse_jpeg(jpeg).each do |mark:0,pos:0,len:0|
      if [0xffd8, 0xfffe, 0xffe0..0xffef].all? { |x| not (x === mark) } then
        data.scan(/.{1,65533}/mn) do |s|
          yield([target, s.size + 2].pack('nn'))
          yield s
        end
        yield(jpeg[pos..-1])
        break
      end
      yield(jpeg[pos, len+2]) unless target === mark
    end
  end

end

if __FILE__ == $0 then
  $stdout.binmode
  $stdin.set_encoding Encoding::ASCII_8BIT
  require 'optparse'
  opts = ARGV.getopts('0123456789abcdefr')
  readmode = opts.delete('r') || ARGV.length == 0
  target = (opts.select{|k,v|v}.keys.map{|x|x.hex}.max || 0x1e) + 0xffe0
  input = ARGV.shift
  jpeg = (input == '-' or input.nil?) ? $stdin.read : File.binread(input)
  if readmode then
    JPGCOM.new.read(jpeg, target) do |data|
      $stdout.write data
    end
  else
    data = $<.set_encoding(Encoding::ASCII_8BIT).read
    JPGCOM.new.write(jpeg, data, target) do |jpeg|
      $stdout.write jpeg
    end
  end
end