update-log.rb

#!/usr/bin/ruby
require 'rubygems'
require 'net/http'
require 'uri'
require 'json'
require 'bindata'
require 'base64'
require 'openssl'
require 'digest/sha1'
require 'pp'

# From http://stackoverflow.com/questions/754494/reading-the-last-n-lines-of-a-file-in-ruby
def tail(path, n)
  file = File.open(path, "r")
  buffer_s = 512
  line_count = 0
  file.seek(0, IO::SEEK_END)

  offset = file.pos # we start at the end

  while line_count <= n && offset > 0
    to_read = if (offset - buffer_s) < 0
                offset
              else
                buffer_s
              end

    file.seek(offset-to_read)
    data = file.read(to_read)

    data.reverse.each_char do |c|
      if line_count > n
        offset += 1
        break
      end
      offset -= 1
      if c == "\n"
        line_count += 1
      end
    end
  end

  file.seek(offset)
  data = file.read
end

LOGSERVER = ARGV[0] 
CERTPATH = '/media/ephemeral0/ctlogs/'
CHAINPATH = '/media/ephemeral0/ctchains/'

class ASN1Cert < BinData::Primitive
	endian :big
	uint24	:len, :value => lambda {data.length}
	string	:data, :read_length => :len

	def get;   self.data; end
	def set(v) self.data = v; end

	def certificate; OpenSSL::X509::Certificate.new(self.data); end
	def to_text; self.certificate.to_text; end
	def to_pem; self.certificate.to_pem; end
	def to_der; self.certificate.to_der; end
end

class PreCert < BinData::Record
	endian :big
	string	:issuer_key_hash, :read_length => 32
	ASN1Cert :tbs_certificate
end

class MerkleTreeLeaf < BinData::Record
	hide :zero, :extensions_len
	endian :big
	uint8 :version, :assert => 0 #v1
	uint8 :leaf_type, :assert => 0 #timestamped_entry
	uint64 :timestamp #Unix timestamp in ms
	uint16 :entry_type
	choice :signed_entry, :selection => :entry_type do
		ASN1Cert 0
		PreCert 1
	end
	uint16 :extensions_len, :assert => 0
	string :extensions, :read_length => :extensions_len
	count_bytes_remaining :zero
        virtual :assert => lambda { zero == 0 }

	def self.read_base64(b64)
		self.read(Base64.decode64(b64))
	end

	def raw_certificate
		if entry_type == 0
			signed_entry
		elsif entry_type == 1
			signed_entry.tbs_certificate
		end
	end
	def certificate
		if entry_type == 0
			signed_entry.certificate
		elsif entry_type == 1
			signed_entry.tbs_certificate.certificate
		end
	end
end

class ChainEntry < BinData::Record
	hide :zero
	endian :big
	uint24	:chain_len
	array :chain, :type => ASN1Cert, :read_until => lambda{array.num_bytes == chain_len}
	count_bytes_remaining :zero
        virtual :assert => lambda { zero == 0 }

	def self.read_base64(b64)
		if b64 == "AAAA"
			# Empty
			return ChainEntry.new
		end
		self.read(Base64.decode64(b64))
	end
end

class PrecertChainEntry < BinData::Record
	endian :big
	ASN1Cert :pre_certificate
	ChainEntry :precertificate_chain

	def self.read_base64(b64)
		self.read(Base64.decode64(b64))
	end
end	

class CT
	def initialize(log = LOGSERVER)
		@log = URI.parse(log + "/").normalize
	end

	def get_sth	
		url = @log + "ct/v1/get-sth"
		_call url
	end

	def get_entries(min = 0, max = 0)
		url = @log + "ct/v1/get-entries"
		url.query = qstr({:start => min, :end => max})
		j = _call(url)
		if !j.has_key? "entries"
			$stderr.puts j
		end
		j["entries"]
	end

	private
	def qstr(h)
		h.map{|k, v| "#{k}=#{v}"}.join("&")
	end

	def _call(url)
		resp = Net::HTTP.get_response url
		if resp.code != '200'
			$stderr.puts(resp.body)
			$stderr.puts "#{resp.code.class}: #{resp.code}"
		end
		JSON.parse(resp.body)
	end
end

CERTFILE=File.join(CERTPATH, ARGV[1])
if (File.size? CERTFILE).nil?
	last_entry = 0
else
	last_entry = tail(CERTFILE, 1).split(",")[0].to_i
end
count = last_entry

outfile = File.open(CERTFILE, "a+")

ct = CT.new
tree_size = ct.get_sth["tree_size"]
$stderr.puts ARGV[0]
$stderr.puts "Total entries: #{tree_size}"
# Entry is zero based, so add one of count
$stderr.puts "Downloaded entries: #{count + 1}"

# Entries sequence is zero based, so max entry is one less than size
max_entry = tree_size - 1
loop do
	max = [max_entry, count + 2000].min
	e = ct.get_entries(count, max)
	if e.nil?
		$stderr.puts "Got nil response for #{count}. Stopping."
		break
	end
	e.each do |x|
		if x.keys.length > 2
		$stderr.puts x.keys.join(" ") 
		end
		$stderr.puts "#{count}" if (count % 5000 == 0)
		leaf = nil
		begin
			leaf = MerkleTreeLeaf.read_base64(x["leaf_input"])
			outfile.puts(count.to_s + "," + leaf.entry_type.to_s + "," + Base64.strict_encode64(leaf.raw_certificate))
			count += 1
		rescue => e
			puts "Bad record #{count}: #{e.class} (#{e.message})"
			count += 1
			next
		end
		begin
			if leaf.entry_type == 0
				chain = ChainEntry.read_base64(x["extra_data"])
			else 
				prechain = PrecertChainEntry.read_base64(x["extra_data"])
				chain = prechain.precertificate_chain
			end
			chain.chain.each do |cert|
				der = cert.to_der
				hash = Digest::SHA1.hexdigest(der)
				if !File.exist? (CHAINPATH + hash)
					puts "New chain entry"
					File.open(CHAINPATH + hash, 'w') {|f|
						f.write(der)
					}
				end
			end	
		rescue
			puts "Bad chain #{count}"
			next
		end
	end
	break if count >= max_entry 
end
outfile.close