Add a script to test the DF state API v2, and factor out IRS response decryption

Gemfile

@@ -143,6 +143,7 @@ group :development do
   gem 'stackprof'
   gem 'memory_profiler'
   gem "letter_opener"
+  gem "faraday"
 end
 
 group :test do

Gemfile.lock

@@ -769,6 +769,7 @@ DEPENDENCIES
   dogapi
   easy_translate
   factory_bot_rails
+  faraday
   fix-db-schema-conflicts
   flamegraph
   flipper

app/services/irs_api_service.rb

@@ -84,34 +84,42 @@ def self.import_federal_data(authorization_code, _state_code)
       return
     end
 
+    decrypted_response = decrypt_response(
+      cert_finder.client_key,
+      Base64.decode64(response.header['SESSION-KEY']),
+      Base64.decode64(response.header['INITIALIZATION-VECTOR']),
+      Base64.decode64(JSON.parse(response.body)['taxReturn']),
+      Base64.decode64(response.header['AUTHENTICATION-TAG'])
+    )
+
+    decrypted_json = JSON.parse(decrypted_response)
+    decrypted_json['xml'] = Nokogiri::XML(decrypted_json['xml']).to_xml
+
+    decrypted_json
+  end
+
+  def self.decrypt_response(private_key, encrypted_secret, initialization_vector, encrypted_data, authentication_tag = nil)
     decipher = OpenSSL::Cipher.new('aes-256-gcm')
     decipher.decrypt
-    client_key = cert_finder.client_key
-    encrypted_session_key = Base64.decode64(response.header['SESSION-KEY'])
 
     label = ''
     md_oaep = OpenSSL::Digest::SHA256
     md_mgf1 = OpenSSL::Digest::SHA1
 
-    decipher.key = client_key.private_decrypt_oaep(encrypted_session_key, label, md_oaep, md_mgf1)
-    decipher.iv = Base64.decode64(response.header['INITIALIZATION-VECTOR'])
-    encrypted_tax_return_bytes = Base64.decode64(JSON.parse(response.body)['taxReturn'])
+    decipher.key = private_key.private_decrypt_oaep(encrypted_secret, label, md_oaep, md_mgf1)
+    decipher.iv = initialization_vector
 
     if ENV['IRS_API_LOCALHOST']
-      decipher.auth_tag = Base64.decode64(response.header['AUTHENTICATION-TAG'])
+      decipher.auth_tag = authentication_tag
     else
-      char_array = encrypted_tax_return_bytes.unpack("C*")
-      encrypted_tax_return_bytes = char_array[0..-17].pack("C*")
+      char_array = encrypted_data.unpack("C*")
+      encrypted_data = char_array[0..-17].pack("C*")
       auth_tag = char_array.last(16).pack("C*")
 
       decipher.auth_tag = auth_tag
     end
-    plain = decipher.update(encrypted_tax_return_bytes) + decipher.final
-
-    decrypted_json = JSON.parse(plain)
-    decrypted_json['xml'] = Nokogiri::XML(decrypted_json['xml']).to_xml
 
-    decrypted_json
+    decipher.update(encrypted_data) + decipher.final
   end
 
   private

scripts/test_state_api_v2.rb

@@ -0,0 +1,85 @@
+#!/usr/bin/env ruby
+
+require_relative "../config/environment"
+
+class TestStateApiV2 < Thor
+  desc 'access_token <kid> <client_id> <private_key_path>', 'Retrieves an OAuth access token for subsequent requests'
+  def access_token(kid, client_id, private_key_path)
+    say "Retrieving an access token", :green
+
+    host = "https://api.alt.www4.irs.gov"
+    path = "/auth/oauth/v2/token"
+
+    headers = {
+      alg: "RS256",
+      typ: "JWT",
+      kid: kid
+    }
+    body = {
+      aud: host + path,
+      iss: client_id,
+      sub: client_id,
+      exp: Time.now.to_i + 15 * 60,
+      jti: SecureRandom.uuid
+    }
+    private_key = OpenSSL::PKey::RSA.new(File.binread(private_key_path))
+    jwt = JWT.encode(body, private_key, "RS256", headers)
+
+    connection = Faraday.new(
+      url: host,
+      headers: {"Content-Type" => " application/x-www-form-urlencoded"}
+    )
+    response = connection.post(path) do |req|
+      req.params = {
+        grant_type: "client_credentials",
+        client_assertion_type: "urn:ietf:params:oauth:client-assertion-type:jwt-bearer",
+        client_assertion: jwt
+      }
+    end
+
+    if JSON.parse(response.body).key?("access_token")
+      say "Retrieved access token: #{JSON.parse(response.body)["access_token"]}", :green
+    else
+      say "Failed to retrieve access token", :red
+      say response.body, :red
+    end
+  end
+
+  desc 'test_connection <access_token> <client_id> <private_key_path>', 'Hits an endpoint meant to test an access token'
+  def test_connection(access_token, client_id, private_key_path)
+    say "Sending request to test connection endpoint", :green
+
+    host = "https://api.alt.www4.irs.gov"
+    path = "/direct-file/state-api/v2/test/connection"
+
+    connection = Faraday.new(
+      url: host,
+      headers: {
+        "Content-Type" => " application/json",
+        "Authorization" => "Bearer #{access_token}",
+        "enterpriseBusCorrelationId" => "#{SecureRandom.uuid}:DFS00:#{client_id[-12..]}:T"
+      }
+    )
+    response = connection.post(path)
+    response_json = JSON.parse(response.body)
+
+    unless response_json.key?("data")
+      say "Received an invalid response", :red
+      say response.body, :red
+      exit
+    end
+
+    decrypted_response = IrsApiService.decrypt_response(
+      OpenSSL::PKey::RSA.new(File.binread(private_key_path)),
+      Base64.decode64(response_json["decryptionInputs"]["encryptedSecret"]),
+      Base64.decode64(response_json["decryptionInputs"]["initializationVector"]),
+      Base64.decode64(response_json["data"]),
+      Base64.decode64(response_json["decryptionInputs"]["authenticationTag"])
+    )
+
+    say "Received a valid response", :green
+    say JSON.parse(decrypted_response), :green
+  end
+end
+
+TestStateApiV2.start