# -*- coding: binary -*-
require 'msf/core'
require 'snmp'

module Msf

###
#
# This module exposes methods for querying a SNMP service
#
###
module Exploit::Remote::SNMPClient


  include Exploit::Remote::Udp


  #
  # Creates an instance of a SNMP exploit module.
  #
  def initialize(info = {})
    super

    # Register the options that all SNMP exploits may make use of.
    register_options(
      [
        Opt::RHOST,
        Opt::RPORT(161),
        OptString.new('COMMUNITY', [ true, 'SNMP Community String', 'public' ]),
        OptString.new('VERSION', [ true, 'SNMP Version <1/2c>', '1' ]),
        OptInt.new('TIMEOUT', [ true, 'SNMP Timeout', 1 ]),
        OptInt.new('RETRIES', [ true, 'SNMP Retries', 1 ])
      ], Msf::Exploit::Remote::SNMPClient)
  end

  #
  # This method wraps the snmp library and passes in
  # the Rex UDP socket
  #
  def connect_snmp(global=true, opts={})
    s = connect_udp(false, opts)

    version = :SNMPv1 if datastore['VERSION'] == '1'
    version = :SNMPv2c if datastore['VERSION'] == '2c'

    snmp = ::SNMP::Manager.new(
      :Host => opts['PeerHost'] || rhost,
      :Port => opts['PeerPort'] || rport,
      :Community => datastore['COMMUNITY'],
      :Version => version,
      :Timeout => datastore['TIMEOUT'],
      :Retries => datastore['RETRIES'],
      :Transport => SNMP::RexUDPTransport,
      :Socket => s
    )

    @snmp = snmp if global
    snmp
  end

  def disconnect_snmp
    @snmp.close if @snmp
    @snmp = nil
  end

  def snmp
    @snmp
  end

  def community
    datastore['COMMUNITY'] || 'public'
  end

  def timeout
    datastore['TIMEOUT'] || 1
  end

  def retries
    datastore['RETRIES'] || 1
  end

protected

  attr_accessor :snmp
end
end
