CoCalc provides the best real-time collaborative environment for Jupyter Notebooks, LaTeX documents, and SageMath, scalable from individual users to large groups and classes!

1
require 'spec_helper'
2

3
load Metasploit::Framework.root.join('tools/exploit/virustotal.rb').to_path
4

5
require 'msfenv'
6
require 'digest/sha2'
7

8
RSpec.describe VirusTotalUtility do
9

10
  context "Classes" do
11
    let(:api_key) do
12
      'FAKE_API_KEY'
13
    end
14

15
    let(:filename) do
16
      'MALWARE.EXE'
17
    end
18

19
    let(:malware_data) do
20
      'DATA'
21
    end
22

23
    describe VirusTotalUtility::ToolConfig do
24
      context "Class methods" do
25

26
        let(:tool_config) do
27
          VirusTotalUtility::ToolConfig.new
28
        end
29

30
        context ".Initializer" do
31
          it "should init the config file path as Metasploit's default config path" do
32
            expect(tool_config.instance_variable_get(:@config_file)).to eq(Msf::Config.config_file)
33
          end
34

35
          it "should init the group name as 'VirusTotal'" do
36
            expect(tool_config.instance_variable_get(:@group_name)).to eq('VirusTotal')
37
          end
38
        end
39
      end
40
    end
41

42
    describe VirusTotalUtility::VirusTotal do
43
      context "Class methods" do
44

45
        let(:malware_sha256) do
46
          Digest::SHA256.hexdigest(malware_data)
47
        end
48

49
        let(:sample) do
50
          {
51
            'filename' => filename,
52
            'data'     => malware_data,
53
            'sha256'   => malware_sha256
54
          }
55
        end
56

57
        let(:boundary) do
58
          'THEREAREMANYLIKEITBUTTHISISMYDATA'
59
        end
60

61
        let(:scan_sample_opts) do
62
          opts = {
63
            'boundary' => boundary,
64
            'api_key'  => api_key,
65
            'filename' => filename,
66
            'data'     => malware_data
67
          }
68

69
          return opts
70
        end
71

72
        let(:retrieve_report_opts) do
73
          opts = {
74
            'uri'       => '/vtapi/v2/file/report',
75
            'method'    => 'POST',
76
            'vhost'     => 'www.virustotal.com',
77
            'vars_post' => {
78
              'apikey'   => api_key,
79
              'resource' => malware_sha256
80
            }
81
          }
82

83
          return opts
84
        end
85

86
        let(:vt) do
87
          file = double(File, read: malware_data)
88
          allow(File).to receive(:open).with(filename, 'rb') {|&block| block.yield file}
89
          VirusTotalUtility::VirusTotal.new({'api_key'=>api_key, 'sample'=>filename})
90
        end
91

92
        context ".Initializer" do
93
          it "should have an API key" do
94
            expect(vt.instance_variable_get(:@api_key)).to eq(api_key)
95
          end
96

97
          it "should have a checksum for the malware sample" do
98
            expect(vt.instance_variable_get(:@sample_info)['sha256']).to eq(malware_sha256)
99
          end
100
        end
101

102
        context "._load_sample" do
103
          it "should contain sample info including data, filename, and sha256" do
104
            expect(vt.send(:_load_sample, filename)).to eq(sample)
105
          end
106
        end
107

108
        context ".scan_sample" do
109
          it "should return with data" do
110
            expect(vt).to receive(:_execute_request).and_return('')
111
            expect(vt.scan_sample).to eq('')
112
          end
113
        end
114

115
        context ".retrieve_report" do
116
          it "should return with data" do
117
            expect(vt).to receive(:_execute_request).and_return('')
118
            expect(vt.retrieve_report).to eq('')
119
          end
120
        end
121

122
        context "._execute_request" do
123
          it "should return status code 204" do
124
            res = double(Rex::Proto::Http::Response)
125
            expect(res).to receive(:code).and_return(204)
126
            expect(vt).to receive(:send_request_cgi).with(scan_sample_opts).and_return(res)
127
            expect { vt.send(:_execute_request, scan_sample_opts) }.to raise_error(RuntimeError)
128
          end
129

130
          it "should return status code 403" do
131
            res = double(Rex::Proto::Http::Response)
132
            expect(res).to receive(:code).and_return(403)
133
            expect(vt).to receive(:send_request_cgi).with(scan_sample_opts).and_return(res)
134
            expect { vt.send(:_execute_request, scan_sample_opts) }.to raise_error(RuntimeError)
135
          end
136
        end
137

138
        context "._create_upload_data" do
139

140
          let(:form_opts) do
141
            {
142
              'boundary' => boundary,
143
              'api_key'  => api_key,
144
              'filename' => filename,
145
              'data'     => malware_data
146
            }
147
          end
148

149
          before(:example) do
150
            @upload_data = vt.send(:_create_upload_data, form_opts)
151
          end
152

153
          it "should create form-data with a boundary" do
154
            expect(@upload_data).to match(/#{boundary}/)
155
          end
156

157
          it "should create form-data with the API key" do
158
            expect(@upload_data).to match(/#{api_key}/)
159
          end
160

161
          it "should create form-data with the malware filename" do
162
            expect(@upload_data).to match(/#{filename}/)
163
          end
164

165
          it "should create form-data with the malware data" do
166
            expect(@upload_data).to match(/#{malware_data}/)
167
          end
168
        end
169
      end
170
    end
171

172

173
    describe VirusTotalUtility::Driver do
174
      before do
175
        $stdin = StringIO.new("Y\n")
176
      end
177

178
      after do
179
        $stdin = STDIN
180
      end
181

182
      let(:driver) do
183
        argv = "-k #{api_key} -f #{filename}".split
184
        options = {
185
          'samples' => filename,
186
          'api_key' => api_key,
187
          'delay'   => 60
188
        }
189

190
        expect(VirusTotalUtility::OptsConsole).to receive(:parse).with(anything).and_return(options)
191

192
        tool_config = instance_double(
193
          VirusTotalUtility::ToolConfig,
194
          has_privacy_waiver?: true,
195
          load_api_key: api_key,
196
          save_api_key: nil,
197
          save_privacy_waiver: nil
198
        )
199

200
        expect(VirusTotalUtility::ToolConfig).to receive(:new).and_return(tool_config)
201

202
        d = nil
203

204
        get_stdout {
205
          d = VirusTotalUtility::Driver.new
206
        }
207

208
        d
209
      end
210

211
      context ".Class methods" do
212

213
        context ".initialize" do
214
          it "should return a Driver object" do
215
            expect(driver.class).to eq(VirusTotalUtility::Driver)
216
          end
217
        end
218

219
        context ".ask_privacy" do
220
          it "should have a link of VirusTotal's terms of service" do
221
            tos = 'https://www.virustotal.com/en/about/terms-of-service'
222
            out = get_stdout { driver.ack_privacy }
223
            expect(out).to match(/#{tos}/)
224
          end
225
        end
226

227
        context ".generate_report" do
228
          it "should show a report" do
229
            res = {
230
              "scans" => {
231
                "Bkav" => { "detected" => false, "version" => "1.3.0.4613", "result" => nil, "update" => "20140107" }
232
              },
233
              "response_code" => 1
234
            }
235

236
            out = get_stdout { driver.generate_report(res, filename) }
237
            expect(out).to match(/#{res['scans']['Bkav']['version']}/)
238
          end
239
        end
240
      end
241
    end
242
  end
243
end
244

245