require 'bindata'

# Implementation of the BitCoin wire protocol, written using bindata.
# Reference: https://en.bitcoin.it/wiki/Protocol_specification
# 
# @author Nick Thomas <nick@lupine.me.uk
module BtcWireProto
  CURRENT_VERSION = 32100
  # Comprehensive list of known networks. The hex values are what you see in
  # MessageHdr#magic and the symbols are their known friendly names.
  NETWORKS = {
    :main      => 0xF9BEB4D9,
    0xF9BEB4D9 => :main,

    :testnet   => 0xFABFB5DA,
    0xFABFB5DA => :testnet
  }
  
  # Comprehensive list of known inventory vector types.
  INV_VEC_TYPES = {
    0 => :error,
    :error => 0,
    
    1 => :msg_tx,
    :msg_tx => 1,
    
    2 => :msg_block,
    :msg_block => 2
  }
  
  # Used in Reply messages
  REPLY_CODES = {
    0 => :success,
    :success => 0,
    
    1 => :wallet_error,
    :wallet_error => 1,
    
    2 => :denied,
    :denied => 2
  }
  
  # Only Alert messages signed by this key are valid.
  # This is an ECDSA public key (FIXME: in what format?)
  ALERT_PUBKEY = "04fc9702847840aaf195de8442ebecedf5b095cdbb9bc716bda9110971b" +
                 "28a49e0ead8564ff0db22209e0374782c093bb899692d524e9d6a6956e7" +
                 "c5ecbcd68284"
  
  ## Components of payloads ##
end
  
  # Bitmask advertising various capabilities of the node.
  # @author Nick Thomas <nick@lupine.me.uk>  
  class BtcWireProto::ServicesMask < BinData::Record
    endian :little
    
    bit62 :undefined_top
    bit1  :node_network
    bit1  :undefined_bottom
  end

  # Structure holding an IP address and port in a slightly unusual format.
  # This one is big-endian - everything else is little-endian.
  #
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::NetAddr < BinData::Record
    endian :big
    services_mask :services
    uint128 :ip # IPv6 address. IPv4 addresses given as IPv6-mapped IPv4
    uint16 :port
  end
  
  # Like a NetAddr but with a timestamp to boot.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::TimestampedNetAddr < BinData::Record
    endian :little

    uint32   :timestamp
    net_addr :net_addr
  end
  
  # Variable-length integer. This is slightly scary.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::VarInt < BinData::BasePrimitive

    def value_to_binary_string(val)
      val = val.to_i

      if val < -0xffffffffffffffff # unrepresentable
        ""
      elsif val < 0  # 64-bit negative integer
        top_32 = (val & 0xffffffff00000000) >> 32
        btm_32 = val & 0x00000000ffffffff
        [0xff, top_32, btm_32].pack("CVV") 
      elsif val <= 0xfc # 8-bit (almost) positive integer
        [val].pack("C")
      elsif val <= 0xffff # 16-bit positive integer
        [0xfd, val].pack("Cv")
      elsif val <= 0xffffffff # 32-bit positive integer
        [0xfe, val].pack("CV")
      else  # We can't represent this, whatever it is
        ""
      end 
    end
    BinData::RegisteredClasses.register("var_int", BtcWireProto::VarInt)
    
    def read_and_return_value(io)
      return nil if io.length < 1

      magic = read_uint8(io)
      if magic <= 0xfc # 8-bit (almost) positive integer
        magic
      elsif magic == 0xfd # 16-bit positive integer
        read_uint16(io)
      elsif magic == 0xfe # 32-bit positive integer
        read_uint32(io)
      elsif magic == 0xff # 64-bit negative integer  
        -(read_uint64(io))
      end
    end
    
    def sensible_default
      0
    end
  end

  # Variable-length pascal string with a variable-length int specifying the
  # length. I kid you not.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::VarStr < BinData::Primitive
    endian :little

    var_int :len, :value => lambda { data.length }
    string  :data, :read_length => :len
    
    def get    ; self.data     ; end
    def set(v) ; self.data = v ; end
  end
  
  class BtcWireProto::InventoryVector < BinData::Record
    endian :little

    uint32 :type # For values, see INV_VEC_TYPES
    string :hash, :length => 32
  end
  
  # Simple class wrapping raw SHA256 data. Might have utility methods later.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Sha256 < BinData::Record
    uint256 :data # Raw SHA256 data
  end
  BtcWireProto::SHA256 = BtcWireProto::Sha256

  class BtcWireProto::TransactionIn < BinData::Record
    struct :previous_output do
      sha256 :hash
      uint32 :index
    end
    var_str :signature_script # Script for confirming transaction authorisation
    uint32 :sequence # Version of this record. 
  end
  
  class BtcWireProto::TranactionOut < BinData::Record
    uint64 :value
    var_str :pk_script # Script containing conditions to claim to transaction
  end

  # Header for a block. 
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::BlockHdr < BinData::Record
    endian :little

    uint32  :version
    sha256  :prev_block
    sha256  :merkle_root
    uint32  :timestamp
    uint32  :difficulty
    uint32  :nonce
    var_int :txn_count       
  end

  ## Payloads ##

  # Payload for a version message
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Version < BinData::Record
    endian :little

    uint32        :version
    services_mask :services
    uint64        :timestamp
    net_addr      :addr_me
    net_addr      :addr_you,     :only_if => lambda { version >= 106 }
    uint64        :nonce,        :only_if => lambda { version >= 106 }
    var_str       :sub_version,  :only_if => lambda { version >= 106 }
    uint32        :start_height, :only_if => lambda { version >= 209 }
  end
  
  # Payload for an addr message in versions earlier than 31402. These are
  # used to get a list of peers to interact with.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::AddrPre31402 < BinData::Record
    endian :little
    
    var_int :count
    array :addrs, :type => :net_addr,
      :read_until => lambda { index == count - 1 }

  end
  
  # Payload for an addr message in versions later than 31402. A timestamp was
  # added to the list of addresses, but otherwise it's the same as AddrPre31402
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::AddrFrom31402 < BinData::Record
    endian :little
    
    var_int :count
    array :timestamped_addrs, :type => :timestamped_net_addr,
      :read_until => lambda { index == count - 1 }
  end
  
  # Payload for a getdata or inv message. This lets the peer advertise the
  # various objects it has knowledge of.
  # @author Nick Thomas <nick@lupine.me.uk>  
  class BtcWireProto::Inventory < BinData::Record
    endian :little

    var_int :count
    array :items, :type => :inventory_vector,
      :read_until => lambda { index == count - 1 }
  end
  
  # Payload for a getblocks or getheaders message. Specifies a set of blocks
  # that the sender wants details of.
  # @author Nick thomas <nick@lupine.me.uk>
  class BtcWireProto::BlockSpec < BinData::Record
    endian :little
    
    uint32  :version
    var_int :start_count
    array   :hash_start, :type => :sha256,
            :read_until => lambda { index == start_count - 1 }    
    # Hash of the last desired block, or 0 to get as many as possible (max: 500)
    sha256  :hash_stop, :length => 32
  end
  
  # A transaction. This contains a number of transactions 'in', and 'out'.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Transaction < BinData::Record
    endian :little
    
    uint32  :version
    var_int :tx_in_count
    array   :transactions_in, :type => :transaction_in,
            :read_until => lambda { index == tx_in_count - 1 }
    var_int :tx_out_count
    array   :transactions_out, :type => :transaction_out,
            :read_until => lambda { index == tx_out_count - 1 }
    uint32  :lock_time
  end
    
  # Details about a particular block. Returned in response to a block request
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Block < BinData::Record
    endian :little
    
    block_hdr :header
    array     :txns, :type => :transaction,
              :read_until => lambda { index == header.txn_count - 1 }
  end
  
  # Headers payloads are returned in response to a getheaders request.
  # Limit of 2,000 entries per message.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Headers < BinData::Record
    endian :little
    
    var_int :count
    array   :block_hdrs, :type => :block_hdr,
            :read_until => lambda { index == count - 1 }
    
  end
  
  # For now, we don't support CheckOrder requests at all. Protocol documentation
  # is lacking! FIXME
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::CheckOrder < BinData::Record
    endian :little
  end
  
  # We don't support SubmitOrder requests either. Receiving either of these will
  # actually break the stream, since we don't even know how long they are. FIXME
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::SubmitOrder < BinData::Record
    endian :little
  end
  
  # Used as a response to a CheckOrder request. 
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Reply < BinData::Record
    endian :little
    
    uint32 :reply # See REPLYCODES for possible values
  end
  
  # A message sent using the p2p network. Signed by a key so you can tell who
  # sent it - if it's signed by a particular key, then we should apparently
  # show the message to the user and cease operation until further notice. Fun!
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Alert < BinData::Record
    endian :little
    
    var_str :message
    var_str :signature
  end

  ## Top-level message format ##
  
  # Found at the start of all Bitcoin messages.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::MessageHdr < BinData::Record
    endian :little
    uint32 :magic
    string :command, :length => 12
    uint32 :payload_len
    uint32 :checksum, :onlyif => :has_checksum?
    
    protected
    
    # version and verack messages don't have a checksum. The rest do.
    # @return[Boolean] does this message header have a checksum field or not?
    def has_checksum?
      command != "version" && command != "verack"
    end 
  end

  # Everything on the wire is a Message.
  # @author Nick Thomas <nick@lupine.me.uk>
  class BtcWireProto::Message < BinData::Record
  
    # @param[Fixnum,nil] version The protocol version. Setting this affects
    # the layout of various fields.
    def initialize(version = nil)
      @version = version || BtcWireProto::CURRENT_VERSION
    end

    message_hdr :header
    
    choice :payload, :selection => :payload_choice do
      version         "version"
      addr_pre_31402  "addr_pre_31402"
      addr_from_31402 "addr_from_31402"
      inventory       "inv"
      inventory       "getdata"
      block_spec      "getblocks"
      block_spec      "getheaders"
      transaction     "tx"
      block           "block"
      headers         "headers"
      checkorder      "checkorder"
      submitorder     "submitorder"
      alert           "alert"
    end

    protected

    # Works out what the payload looks like based on the MessageHdr struct
    # and (potentially) the version
    def payload_choice
      return header.command if %w{
        version inv getdata getblocks getheaders tx block headers alert
      }.include?(header.command)
      
      # We can't parse these yet, and so we don't know where in the stream the
      # next message starts. So all we can do is throw an error
      raise NotImplementedError.new(
        "Received unsupported command #{header.command}"
      ) if %w|checkorder submitorder|.include?(header.command)
      
      # These commands don't have any payloads
      return nil if %w|verack getaddr ping|.include?(header.command)
      
      # Payload has two forms, depending on protocol version. Ugh.
      return (@version < 31402 ? "addr_pre_31402" : "addr_from_31402") if
        header.command == "addr"
        
    end

  end
#end