# -*- coding: binary -*-
require 'rex/proto/drda'

module Rex
module Proto
module DRDA

class Error < StandardError; end
class RespError < Error; end

# See:
# http://publib.boulder.ibm.com/infocenter/dzichelp/v2r2/index.jsp?topic=/com.ibm.db29.doc.drda/db2z_excsat.htm
class MGRLVLLS_PARAM < Struct.new(:length, :codepoint, :payload)
  def initialize(args={})
    self[:codepoint] = Constants::MGRLVLLS
    self[:payload] = "\x14\x03\x00\x0a\x24\x07\x00\x0a" +
      "\x14\x74\x00\x05\x24\x0f\x00\x08" +
      "\x14\x40\x00\x09\x1c\x08\x04\xb8"
    self[:length] = self[:payload].to_s.size+4
  end
  def to_s
    self.to_a.pack("nna*")
  end
end

# Currently, only takes a MGRLVLLS param. Extend the struct
# when more parameters are defined.
class EXCSAT_DDM < Struct.new(:length, :magic, :format, :correlid, :length2,
   :codepoint, :mgrlvlls)

  def initialize(args={})
    self[:magic] = 0xd0
    self[:format] = 0x41
    self[:correlid] = 1
    self[:codepoint] = Constants::EXCSAT
    self[:mgrlvlls] = args[:mgrlvlls] || MGRLVLLS_PARAM.new.to_s
    self[:length] = (10 + self[:mgrlvlls].to_s.size)
    self[:length2] = self[:length]-6
  end

  def to_s
    packstr = "nCCnnn"
    packstr += "a*"  # Pack smarter as more params are added.
    self.to_a.pack(packstr)
  end
end

# See http://publib.boulder.ibm.com/infocenter/dzichelp/v2r2/index.jsp?topic=/com.ibm.db29.doc.drda/db2z_accsec.htm
# for all sorts of info about SECMEC.
class SECMEC_PARAM < Struct.new(:length, :codepoint, :payload)
  def initialize(args={})
    self[:length] = 6
    self[:codepoint] = Constants::SECMEC
    self[:payload] = 3 # Plaintext username and password.
  end
  def to_s
    self.to_a.pack("nnn")
  end
end

# Relational Database name parameter.
class RDBNAM_PARAM < Struct.new(:length, :codepoint, :payload)
  def initialize(args={})
    self[:length] = 22 # Since the database name is padded out.
    self[:codepoint] = Constants::RDBNAM
    self[:payload] = encode(args[:payload].to_s)
  end

  def encode(str)
    Rex::Text.to_ebcdic([str].pack("A18"))
  end

  def payload=(str)
    self[:payload] = encode(str.to_s)
  end

  def to_s
    self.to_a.pack("nna18")
  end

end

# The ACCSEC DDM is responsible for picking the security mechanism (SECMEC)
# which, in our case, will always be plain text username and password. It
# also sets the relational database name (RDBNAM), if specified. You need
# one to login, but not to probe.
class ACCSEC_DDM < Struct.new(:length, :magic, :format, :correlid, :length2,
   :codepoint, :secmec, :rdbnam)
  def initialize(args={})
    self[:magic] = 0xd0
    self[:format] = args[:format] || 0x01
    self[:correlid] = 2
    self[:codepoint] = Constants::ACCSEC
    self[:secmec] = SECMEC_PARAM.new.to_s
    if args[:dbname] # Include a database name if we're given one.
      self[:rdbnam] = RDBNAM_PARAM.new(:payload => args[:dbname]).to_s
    end
    self[:length] =  10 + self[:secmec].to_s.size + self[:rdbnam].to_s.size
    self[:length2] = self[:length]-6
  end
  def dbname=(str)
    self[:rdbnam] = RDBNAM_PARAM.new(:payload => args[:dbname]).to_s
  end
  def to_s
    packstr = "nCCnnna6"
    packstr += "a22" if self[:rdbnam]
    self.to_a.pack(packstr)
  end
end

class DDM_PARAM < Struct.new(:length, :codepoint, :payload)

  def read(str="")
    raise DRDA::Error, "Input isn't a String." if !str.kind_of? String
    raise DRDA::RespError, "DDM_PARAM is too short" if str.size < 4
    (self[:length], self[:codepoint]) =
      str.unpack("nn")
    raise DRDA::RespError, "DDM_PARAM Length is too short" if self[:length] < 4
    rest = str[4,self[:length]-4] # If it's negative or whatever, it'll end up as "".
    self[:payload] = rest.to_s[0,self[:length]-4]
    return self
  end

  def to_s
    self.to_a.pack("nna*")
  end

end

class BASIC_DDM < Struct.new(:length, :magic, :format, :correlid,
  :length2, :codepoint, :payload)
  def initialize
    self[:payload] = []
  end

  def read(str="")
    self[:payload].clear
    raise DRDA::Error, "Input isn't a String." if !str.kind_of? String
    raise DRDA::RespError, "Response is too short." if str.size < 10
    (self[:length],self[:magic],self[:format],
     self[:correlid],self[:length2],self[:codepoint]) =
     str.unpack("nCCnnn")
    sanity_check
    rest = str[10,self[:length2]-4]
    i = 0
    while (i < rest.size)
      if self[:codepoint] == Constants::SQLCARD # These aren't DDM's.
        this_param = rest[i,self[:length]-10]
      else
        this_param = DDM_PARAM.new.read(rest[i,rest.size])
      end
      self[:payload] << this_param
      i += this_param.to_s.size
    end
    return self
  end

  # Just a quick test.
  def sanity_check
    if self[:length] < 10
      raise DRDA::RespError, "DDM Length is too short."
    elsif self[:length2] < 4
      raise DRDA::RespError, "DDM Length2 is too short."
    elsif self[:length]-6 != self[:length2]
      raise DRDA::RespError, "Codepoint: 0x#{self[:codepoint].to_s(16)} DDM Length2 (0x#{self[:length2].to_s(16)}) isn't six less than Length (0x#{self[:length].to_s(16)})"
    end
  end

  def to_s
    self.to_a.pack("nCCnnn") + self[:payload].map {|x| x.to_s}.join
  end

end

class SERVER_PACKET < Array

  def read(str="")
    raise DRDA::Error, "Input isn't a String." if !str.kind_of? String
    self.clear
    i = 0
    while(i < str.size)
      this_ddm = BASIC_DDM.new.read(str[i,str.size])
      self << this_ddm
      i += this_ddm.to_s.size
    end
    return self
  end

  def to_s; self.join; end
  def sz; self.to_s.size; end

end

class PASSWORD_PARAM < Struct.new(:length, :codepoint, :payload)
  def initialize(args={})
    self[:codepoint] = Constants::PASSWORD
    self[:payload] = Rex::Text.to_ebcdic(args[:payload].to_s)
    self[:length] = self[:payload].size + 4
  end
  def encode(str)
    Rex::Text.to_ebcdic(str)
  end
  def to_s
    self.to_a.pack("nna*")
  end
end

class USERID_PARAM < Struct.new(:length, :codepoint, :payload)
  def initialize(args={})
    self[:codepoint] = Constants::USERID
    self[:payload] = Rex::Text.to_ebcdic(args[:payload].to_s)
    self[:length] = self[:payload].size + 4
  end
  def encode(str)
    Rex::Text.to_ebcdic(str)
  end
  def to_s
    self.to_a.pack("nna*")
  end
end

class SECCHK_DDM < Struct.new(:length, :magic, :format, :correlid, :length2,
  :codepoint, :secmec, :rdbnam, :password, :userid)
  def initialize(args={}) # Takes :dbname, :dbpass, :dbuser
    self[:magic] = 0xd0
    self[:format] = 0x01
    self[:correlid] = 2
    self[:codepoint] = Constants::SECCHK
    self[:secmec] = SECMEC_PARAM.new.to_s
    if args[:dbname] # Include a database name if we're given one.
      self[:rdbnam] = RDBNAM_PARAM.new(:payload => args[:dbname]).to_s
    end
    self[:password] = PASSWORD_PARAM.new(:payload => args[:dbpass]).to_s
    self[:userid] = USERID_PARAM.new(:payload => args[:dbuser]).to_s
    self[:length] = ( 10 + self[:secmec].to_s.size + self[:rdbnam].to_s.size +
           self[:password].to_s.size + self[:userid].to_s.size )
    self[:length2] = self[:length]-6
  end
  def dbname=(str)
    self[:rdbnam] = RDBNAM_PARAM.new(:payload => args[:dbname]).to_s
  end
  def to_s
    packstr = "nCCnnna6"
    packstr += "a22" if self[:rdbnam]
    packstr += "a*a*" # username and password
    self.to_a.pack(packstr)
  end
end

end
end
end