From 7b807d4dce1808f995ea491db9f8dbb5e26de66a Mon Sep 17 00:00:00 2001
From: Jacob Robles <jrobles@rapid7.com>
Date: Fri, 19 Apr 2019 06:28:56 -0500
Subject: [PATCH] Add nuuo client rex and spec

---
 lib/rex/proto/nuuo/client.rb           |  67 +--
 spec/lib/rex/proto/nuuo/client_spec.rb | 537 +++++++++++++++++++++++++
 2 files changed, 577 insertions(+), 27 deletions(-)
 create mode 100644 spec/lib/rex/proto/nuuo/client_spec.rb

diff --git a/lib/rex/proto/nuuo/client.rb b/lib/rex/proto/nuuo/client.rb
index 882ff48974..aeb4af8176 100644
--- a/lib/rex/proto/nuuo/client.rb
+++ b/lib/rex/proto/nuuo/client.rb
@@ -1,5 +1,7 @@
 # -*- coding: binary -*-
 
+require 'rex/proto/nuuo/client_request'
+
 module Rex
 module Proto
 module Nuuo
@@ -56,9 +58,9 @@ class Client
   #
   # @return [Rex::Socket::Tcp]
   # @raise [RuntimeError] if 'tcp' is not requested
-  def connect
-    return connection if connection
-    return  create_tcp_connection if protocol == 'tcp'
+  def connect(temp: false)
+    return connection if connection && !temp
+    return create_tcp_connection(temp: temp) if protocol == 'tcp'
     raise ::RuntimeError, 'Nuuo Client: Unknown transport protocol'
   end
 
@@ -72,25 +74,34 @@ class Client
     self.connection = nil
   end
 
-  def send_recv(req)
-    send_request(req)
-    read_response
+  def send_recv(req, conn=nil)
+    send_request(req, conn)
+    read_response(conn)
   end
 
-  def send_request(req)
-    connect.put(req.to_s)
+  def send_request(req, conn=nil)
+    conn ? conn.put(req.to_s) : connect.put(req.to_s)
   end
 
-  def read_response
-    connection.get_once
+  def read_response(conn=nil)
+    data = conn ? conn.get_once : connection.get_once
+  end
+
+  def user_session_header(opts)
+    val = nil
+    if opts['user_session']
+      val = opts['user_session']
+    elsif self.user_session
+      val = self.user_session
+    end
   end
 
   def request_ping(opts={})
     opts = self.config.merge(opts)
     opts['headers'] ||= {}
     opts['method'] = 'PING'
-
-    opts['headers']['User-Session-No'] = opts['user_session']
+    session = user_session_header(opts)
+    opts['headers']['User-Session-No'] = session if session
 
     ClientRequest.new(opts)
   end
@@ -100,11 +111,11 @@ class Client
     opts['headers'] ||= {}
     opts['method'] = 'SENDLICFILE'
 
+    session = user_session_header(opts)
+    opts['headers']['User-Session-No'] = session if session
+    opts['data'] = '' unless opts['data']
+
     opts['headers']['FileName'] = opts['file_name']
-    opts['headers']['User-Session-No'] = opts['user_session']
-    unless opts['data']
-      opts['data'] = ''
-    end
     opts['headers']['Content-Length'] = opts['data'].length
 
     ClientRequest.new(opts)
@@ -122,7 +133,8 @@ class Client
 
     opts['headers']['FileName'] = opts['file_name']
     opts['headers']['FileType'] = opts['file_type'] || 1
-    opts['headers']['User-Session-No'] = opts['user_session']
+    session = user_session_header(opts)
+    opts['headers']['User-Session-No'] = session if session
 
     ClientRequest.new(opts)
   end
@@ -142,10 +154,11 @@ class Client
 
     opts['headers']['FileName'] = opts['file_name']
     opts['headers']['FileType'] = opts['file_type'] || 1
-    opts['headers']['User-Session-No'] = opts['user_session']
-    unless opts['data']
-      opts['data'] = ''
-    end
+
+    session = user_session_header(opts)
+    opts['headers']['User-Session-No'] = session if session
+
+    opts['data'] = '' unless opts['data']
     opts['headers']['Content-Length'] = opts['data'].length
 
     ClientRequest.new(opts)
@@ -167,7 +180,7 @@ class Client
     # Account for version...
     opts['headers']['Version'] = opts['server_version']
 
-    username = ''
+    username = nil
     if opts['username'] && opts['username'] != ''
       username = opts['username']
     elsif self.username && self.username != ''
@@ -175,7 +188,6 @@ class Client
     end
 
     opts['headers']['Username'] = username
-    opts['username'] = username
 
     password = ''
     if opts['password'] && opts['password'] != ''
@@ -183,9 +195,8 @@ class Client
     elsif self.password && self.password != ''
       password = self.password
     end
-    opts['headers']['Password-Length'] = password.length
-    opts['password'] = password
     opts['data'] = password
+    opts['headers']['Password-Length'] = password.length
 
     # Need to verify if this is needed
     opts['headers']['TimeZone-Length'] = '0'
@@ -216,13 +227,15 @@ class Client
   # Creates a TCP connection using Rex::Socket::Tcp
   #
   # @return [Rex::Socket::Tcp]
-  def create_tcp_connection
-    self.connection = Rex::Socket::Tcp.create(
+  def create_tcp_connection(temp: false)
+    tcp_connection = Rex::Socket::Tcp.create(
       'PeerHost'  => host,
       'PeerPort'  => port.to_i,
       'Context'   => context,
       'Timeout'   => timeout
     )
+    self.connection = tcp_connection unless temp
+    tcp_connection
   end
 
 end
diff --git a/spec/lib/rex/proto/nuuo/client_spec.rb b/spec/lib/rex/proto/nuuo/client_spec.rb
new file mode 100644
index 0000000000..a1db6bae79
--- /dev/null
+++ b/spec/lib/rex/proto/nuuo/client_spec.rb
@@ -0,0 +1,537 @@
+# -*- coding:binary -*-
+require 'rex/proto/nuuo/client'
+
+RSpec.describe Rex::Proto::Nuuo::Client do
+  subject(:client) {
+    described_class.new({
+      protocol: protocol,
+      user_session: client_user_session,
+      username: client_username,
+      password: client_password
+    })
+  }
+  let(:protocol) {'tcp'}
+  let(:client_user_session) {nil}
+  let(:client_username) {nil}
+  let(:client_password) {nil}
+
+  describe '#connect' do
+
+    context 'given udp option when created' do
+      let(:protocol) {'udp'}
+
+      it 'raises an error' do
+        expect{client.connect}.to raise_error(::RuntimeError)
+      end
+    end
+
+    context 'given temp is false' do
+      context 'when there is no connection' do
+        it 'returns a tcp connection' do
+          tcp_connection = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection)
+
+          expect(client.connect).to eq(tcp_connection)
+        end
+
+        it 'saves the tcp connection' do
+          tcp_connection = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection)
+
+          client.connect
+          expect(client.connection).to eq(tcp_connection)
+        end
+      end
+
+      context 'when there is saved connection' do
+        it 'returns the saved tcp connection' do
+          tcp_connection = double('tcp_connection')
+          client.connection = tcp_connection
+
+          expect(client.connect).to eq(tcp_connection)
+        end
+      end
+    end
+
+    context 'given temp is true' do
+      context 'when there is a saved connection' do
+        it 'returns a new connection' do
+          tcp_connection0 = double('tcp_connection')
+          tcp_connection1 = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection1)
+
+          client.connection = tcp_connection0
+          expect(client.connect(temp: true)).to eq(tcp_connection1)
+        end
+
+        it 'does not overwrite existing connection' do
+          tcp_connection0 = double('tcp_connection')
+          tcp_connection1 = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection1)
+
+          client.connection = tcp_connection0
+          client.connect(temp: true)
+          expect(client.connection).to eq(tcp_connection0)
+        end
+      end
+
+      context 'when there is no saved connection' do
+        it 'returns a new connection' do
+          tcp_connection = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection)
+
+          expect(client.connect(temp: true)).to eq(tcp_connection)
+        end
+
+        it 'does not save the connection' do
+          tcp_connection = double('tcp_connection')
+          allow(Rex::Socket::Tcp).to receive(:create).and_return(tcp_connection)
+
+          client.connect(temp: true)
+          expect(client.connection).to be_nil
+        end
+      end
+    end
+
+  end
+
+  describe '#close' do
+    context 'given there is a connection' do
+      it 'calls shutdown on the connection' do
+        tcp_connection = double('tcp_connection')
+        allow(tcp_connection).to receive(:shutdown) {true}
+        allow(tcp_connection).to receive(:closed?) {false}
+        allow(tcp_connection).to receive(:close) {true}
+        client.connection = tcp_connection
+
+        expect(tcp_connection).to receive(:shutdown)
+        client.close
+      end
+
+      it 'calls closed on the connection' do
+        tcp_connection = double('tcp_connection')
+        allow(tcp_connection).to receive(:shutdown) {true}
+        allow(tcp_connection).to receive(:closed?) {false}
+        allow(tcp_connection).to receive(:close) {true}
+        client.connection = tcp_connection
+
+        expect(tcp_connection).to receive(:close)
+        client.close
+      end
+    end
+  end
+
+  describe '#send_recv' do
+    context 'given no connection is passed in' do
+      it 'uses client connection' do
+        tcp_connection = double('tcp_connection')
+        allow(tcp_connection).to receive(:put)
+        allow(tcp_connection).to receive(:get_once)
+        client.connection = tcp_connection
+
+        expect(tcp_connection).to receive(:put)
+        client.send_recv('test')
+      end
+    end
+
+    context 'given a connection is passed in' do
+      it 'uses the passed in connection' do
+        tcp_connection = double('tcp_connection')
+        passed_connection = double('passed_connection')
+        client.connection = tcp_connection
+
+        allow(passed_connection).to receive(:put)
+        allow(passed_connection).to receive(:get_once)
+
+        expect(passed_connection).to receive(:put)
+        client.send_recv('test', passed_connection)
+      end
+    end
+  end
+
+  describe '#request_ping' do
+    subject(:ping_request) {
+      opts = {'user_session' => user_session}
+      client.request_ping(opts)
+    }
+    let(:user_session) {nil}
+
+    it 'returns a PING client request' do
+      expect(ping_request.to_s).to start_with('PING')
+    end
+
+    context 'given a user_session option' do
+      let(:user_session) {'test'}
+
+      context 'when the client does not have a session' do
+        it 'uses the user_session option' do
+          expect(ping_request.to_s).to match('User-Session-No: test')
+        end
+      end
+
+      context 'when the client has a session' do
+        let(:client_user_session) {'client'}
+
+        it 'overrides the client session value' do
+          expect(ping_request.to_s).to match('User-Session-No: test')
+        end
+      end
+    end
+
+
+    context 'given no user_session is provided' do
+      context 'when the client does not have a session' do
+        it 'does not have a User-Session-No header' do
+          expect(ping_request.to_s).to_not match('User-Session-No:')
+        end
+      end
+
+      context 'when the client has a session' do
+        let(:client_user_session) {'client'}
+
+        it 'uses the client session' do
+          expect(ping_request.to_s).to match('User-Session-No: client')
+        end
+      end
+    end
+
+  end
+
+  describe '#request_sendlicfile' do
+    subject(:sendlicfile_request) {
+      opts = {
+        'file_name' => filename,
+        'data' => data
+      }
+      client.request_sendlicfile(opts).to_s
+    }
+    let(:filename) {'TestFile'}
+    let(:data) {'testdata'}
+
+    it 'returns a SENDLICFILE client request' do
+      expect(sendlicfile_request).to start_with('SENDLICFILE')
+    end
+
+    context 'given file_name' do
+      it 'sets the FileName header with the value' do
+        expect(sendlicfile_request).to match("[^\r\n]\r\nFileName: TestFile\r\n")
+      end
+    end
+
+    context 'given no file_name' do
+      let(:filename) {nil}
+
+      it 'creates an empty FileName header' do
+        expect(sendlicfile_request).to match("[^\r\n]\r\nFileName: \r\n")
+      end
+    end
+
+    context 'given data' do
+      it 'sets the body to the data contents' do
+        expect(sendlicfile_request).to end_with("\r\n\r\ntestdata")
+      end
+
+      it 'sets the Content-Length header with data length' do
+        expect(sendlicfile_request).to match("[^\r\n]\r\nContent-Length: 8\r\n")
+      end
+    end
+
+    context 'given no data' do
+      let(:data) {nil}
+      it 'creates an empty body' do
+        expect(sendlicfile_request).to end_with("\r\n\r\n")
+      end
+
+      it 'set Content-Length header to 0' do
+        expect(sendlicfile_request).to match("[^\r\n]\r\nContent-Length: 0\r\n")
+      end
+    end
+  end
+
+  describe '#request_getconfig' do
+    subject(:getconfig_request) {
+      opts = {
+        'file_name' => filename,
+        'file_type' => filetype
+      }
+      client.request_getconfig(opts).to_s
+    }
+    let(:filename) {'TestName'}
+    let(:filetype) {2}
+
+    it 'returns a GETCONFIG client request' do
+      expect(getconfig_request).to start_with('GETCONFIG')
+    end
+
+    context 'given file_name' do
+      it 'sets the FileName header' do
+        expect(getconfig_request).to match("[^\r\n]\r\nFileName: TestName\r\n")
+      end
+    end
+
+    context 'given no file_name' do
+      let(:filename) {nil}
+      it 'creates an empty FileName header' do
+        expect(getconfig_request).to match("[^\r\n]\r\nFileName: \r\n")
+      end
+    end
+
+    context 'given a file_type' do
+      it 'sets the FileType header' do
+        expect(getconfig_request).to match("[^\r\n]\r\nFileType: 2\r\n")
+      end
+    end
+
+    context 'given no file_type' do
+      let(:filetype) {nil}
+      it 'defaults to 1' do
+        expect(getconfig_request).to match("[^\r\n]\r\nFileType: 1\r\n")
+      end
+    end
+  end
+
+  describe '#request_commitconfig' do
+    subject(:commitconfig_request) {
+      opts = {
+        'file_name' => filename,
+        'file_type' => filetype,
+        'data' => data
+      }
+      client.request_commitconfig(opts).to_s
+    }
+    let(:filename) {'TestName'}
+    let(:filetype) {2}
+    let(:data) {'testdata'}
+
+    it 'returns a COMMITCONFIG client request' do
+      expect(commitconfig_request).to start_with('COMMITCONFIG')
+    end
+
+    context 'given file_name' do
+      it 'sets the FileName header' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nFileName: TestName\r\n")
+      end
+    end
+    
+    context 'given no file_name' do
+      let(:filename) {nil}
+
+      it 'creates an empty FileName header' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nFileName: \r\n")
+      end
+    end
+
+    context 'given file_type' do
+      it 'sets the FileType header' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nFileType: 2\r\n")
+      end
+    end
+
+    context 'given no file_type' do
+      let(:filetype) {nil}
+
+      it 'creates an empty FileType header' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nFileType: 1\r\n")
+      end
+    end
+
+    context 'given data' do
+      it 'sets the request body to the data' do
+        expect(commitconfig_request).to end_with("\r\n\r\ntestdata")
+      end
+
+      it 'sets Content-Length to data length' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nContent-Length: 8\r\n")
+      end
+    end
+
+    context 'given no data' do
+      let(:data) {nil}
+
+      it 'creates an empty request body' do
+        expect(commitconfig_request).to end_with("\r\n\r\n")
+      end
+
+      it 'creates Content-Length header with 0' do
+        expect(commitconfig_request).to match("[^\r\n]\r\nContent-Length: 0\r\n")
+      end
+    end
+  end
+
+  describe '#request_userlogin' do
+    subject(:userlogin_request) {
+      opts = {
+        'server_version' => server_version,
+        'username' => username,
+        'password' => password
+      }
+      client.request_userlogin(opts).to_s
+    }
+    let(:server_version) {'1.1.1'}
+    let(:username) {'user'}
+    let(:password) {'pass'}
+
+    it 'returns a USERLOGIN client request' do
+      expect(userlogin_request).to start_with('USERLOGIN')
+    end
+
+    context 'given server_version' do
+      it 'sets Version header with value' do
+        expect(userlogin_request).to match("[^\r\n]\r\nVersion: 1.1.1\r\n")
+      end
+    end
+
+    context 'given no server_version' do
+      let(:server_version) {nil}
+
+      it 'creates an empty Version header' do
+        expect(userlogin_request).to match("[^\r\n]\r\nVersion: \r\n")
+      end
+    end
+
+    context 'when client has username' do
+      let(:client_username) {'client_user'}
+
+      context 'given username' do
+        it 'sets the Username header with opts username' do
+          expect(userlogin_request).to match("[^\r\n]\r\nUsername: user\r\n")
+        end
+      end
+
+      context 'given no username' do
+        let(:username) {nil}
+
+        it 'creates an Username header with client username' do
+          expect(userlogin_request).to match("[^\r\n]\r\nUsername: client_user\r\n")
+        end
+      end
+    end
+
+    context 'when client has no username' do
+      context 'given username' do
+        it 'sets the Username header with value' do
+          expect(userlogin_request).to match("[^\r\n]\r\nUsername: user\r\n")
+        end
+      end
+
+      context 'given no username' do
+        let(:username) {nil}
+
+        it 'creates an empty Username header' do
+          expect(userlogin_request).to match("[^\r\n]\r\nUsername: \r\n")
+        end
+      end
+    end
+
+    context 'when client has password' do
+      let(:client_password) {'client_pass'}
+
+      context 'given password' do
+        it 'sets body with password' do
+          expect(userlogin_request).to end_with("\r\n\r\npass")
+        end
+
+        it 'sets Password-Length header' do
+          expect(userlogin_request).to match("[^\r\n]\r\nPassword-Length: 4\r\n")
+        end
+      end
+
+      context 'given no password' do
+        let(:password) {nil}
+
+        it 'sets body to client password' do
+          expect(userlogin_request).to end_with("\r\n\r\nclient_pass")
+        end
+
+        it 'creates Password-Length with client password length' do
+          expect(userlogin_request).to match("[^\r\n]\r\nPassword-Length: 11\r\n")
+        end
+      end
+    end
+
+    context 'when client has no password' do
+      context 'given password' do
+        it 'sets body with password' do
+          expect(userlogin_request).to end_with("\r\n\r\npass")
+        end
+
+        it 'sets Password-Length header' do
+          expect(userlogin_request).to match("[^\r\n]\r\nPassword-Length: 4\r\n")
+        end
+      end
+
+      context 'given no password' do
+        let(:password) {nil}
+
+        it 'sets empty body' do
+          expect(userlogin_request).to end_with("\r\n\r\n")
+        end
+
+        it 'creates Password-Length with 0' do
+          expect(userlogin_request).to match("[^\r\n]\r\nPassword-Length: 0\r\n")
+        end
+      end
+    end
+
+  end
+
+  describe '#request_getopenalarm' do
+    subject(:getopenalarm_request) {
+      opts = {
+        'device_id' => device_id,
+        'source_server' => source_server,
+        'last_one' => last_one
+      }
+      client.request_getopenalarm(opts).to_s
+    }
+    let(:device_id) {nil}
+    let(:source_server) {nil}
+    let(:last_one) {nil}
+
+    it 'returns a GETOPENALARM client request' do
+      expect(getopenalarm_request).to start_with('GETOPENALARM')
+    end
+
+    context 'given device_id' do
+      let(:device_id) {2}
+
+      it 'sets DeviceID header with value' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nDeviceID: 2\r\n")
+      end
+    end
+
+    context 'given no device_id' do
+      it 'sets DeviceID header to 1' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nDeviceID: 1\r\n")
+      end
+    end
+
+    context 'given source_server' do
+      let(:source_server) {2}
+
+      it 'sets SourceServer header with value' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nSourceServer: 2\r\n")
+      end
+    end
+
+    context 'given no source_server' do
+      it 'set SourceServer header to 1' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nSourceServer: 1\r\n")
+      end
+    end
+
+    context 'given last_one' do
+      let(:last_one) {2}
+
+      it 'sets LastOne header with value' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nLastOne: 2\r\n")
+      end
+    end
+
+    context 'given no last_one' do
+      it 'sets LastOne to 1' do
+        expect(getopenalarm_request).to match("[^\r\n]\r\nLastOne: 1\r\n")
+      end
+    end
+  end
+end