Modify HTTP response (#15)

* Implement response modification

* Test for response modification

* Update README

README.md

@@ -13,8 +13,8 @@ NetShears adds a Request interceptor mechanisms to be able to modify the HTTP/HT
 - [x] Intercept HTTP/HTTPS request endpoint
 - [x] View traffic logs
 - [x] Request obserever
-- [ ] Intercept HTTP/HTTPS response body
-- [ ] Block HTTP requets
+- [x] Intercept HTTP/HTTPS response body
+- [ ] Block HTTP requests
 
 ## How it works
 
@@ -58,6 +58,17 @@ let endpointModifier = RequestEvaluatorModifierEndpoint(redirectedRequest: endpo
 NetShears.shared.modify(modifier: endpointModifier)
 ```
 
+Response Modification
+
+```swift
+let response = HTTPResponseModifyModel(
+    url: "https://example.com/",
+    data: #"{"message": "ok"}"#.data(using: .utf8)!
+)
+let responseModifier = RequestEvaluatorModifierResponse(response: response)
+NetShears.shared.modify(modifier: responseModifier)
+```
+
 # Traffic Monitoring
 
 Make sure to call ```startLogger()``` before showing netwrok traffic logs.

Sources/NetShears/NetShears.swift

@@ -17,7 +17,7 @@ public final class NetShears: NSObject {
     let networkRequestInterceptor = NetworkRequestInterceptor()
 
     lazy var config: NetworkInterceptorConfig = {
-        var savedModifiers = [RequestEvaluatorModifier]().retrieveFromDisk()
+        var savedModifiers = [Modifier]().retrieveFromDisk()
         return NetworkInterceptorConfig(modifiers: savedModifiers)
     }()
 
@@ -57,11 +57,11 @@ public final class NetShears: NSObject {
         checkSwizzling()
     }
 
-    public func modify(modifier: RequestEvaluatorModifier) {
+    public func modify(modifier: Modifier) {
         config.addModifier(modifier: modifier)
     }
 
-    public func modifiedList() -> [RequestEvaluatorModifier] {
+    public func modifiedList() -> [Modifier] {
         return config.modifiers
     }
 

Sources/NetShears/NetworkInterceptorConfig.swift

@@ -27,22 +27,53 @@ public struct HeaderModifyModel: Codable, Equatable {
     }
 }
 
+public struct HTTPResponseModifyModel: Codable, Equatable {
+    public let url: String
+    public let httpMethod: String
+
+    public let statusCode: Int
+    public let data: Data
+    public let httpVersion: String?
+    public let headers: [String: String]
+
+    public init(
+        url: String,
+        data: Data,
+        httpMethod: String = "GET",
+        statusCode: Int = 200,
+        httpVersion: String? = nil,
+        headers: [String : String] = [:]
+    ) {
+        self.url = url
+        self.data = data
+        self.httpMethod = httpMethod
+        self.statusCode = statusCode
+        self.httpVersion = httpVersion
+        self.headers = headers
+    }
+
+    public var response: URLResponse? {
+        guard let url = URL(string: url) else { return nil }
+        return HTTPURLResponse(url: url, statusCode: statusCode, httpVersion: httpVersion, headerFields: headers)
+    }
+}
+
 final class NetworkInterceptorConfig {
-    var modifiers: [RequestEvaluatorModifier] = [] {
+    var modifiers: [Modifier] = [] {
         didSet {
             modifiers.store()
         }
     }
 
-    init(modifiers: [RequestEvaluatorModifier] = []) {
+    init(modifiers: [Modifier] = []) {
         self.modifiers = modifiers
     }
 
-    func addModifier(modifier: RequestEvaluatorModifier) {
+    func addModifier(modifier: Modifier) {
         self.modifiers.append(modifier)
     }
 
-    func getModifiers() -> [RequestEvaluatorModifier] {
+    func getModifiers() -> [Modifier] {
         return self.modifiers
     }
 

Sources/NetShears/RequestModifier/NetShearsModfierProtocol.swift

@@ -7,30 +7,50 @@
 
 import Foundation
 
+public protocol Modifier: RequestEvaluator, RequestModifierStorage, Codable {}
+
 public protocol RequestEvaluator {
     func isActionAllowed(urlRequest: URLRequest) -> Bool
 }
 
-public protocol RequestModifier {
+public protocol RequestModifier: Modifier {
     func modify(request: inout URLRequest)
 }
 
-public protocol RequestEvaluatorModifier : RequestEvaluator, RequestModifier, Codable {
+public protocol RequestActionModifier: Modifier {
+    func modify(client: URLProtocolClient?, urlProtocol: URLProtocol)
+}
+
+public protocol RequestModifierStorage {
     static var storeFileName: String { get }
 }
 
-extension Array where Element == RequestEvaluatorModifier {
-    func store() {
-        let headers: [RequestEvaluatorModifierHeader] = compactMap { $0 as? RequestEvaluatorModifierHeader }
-        let endpoints: [RequestEvaluatorModifierEndpoint] = compactMap { $0 as? RequestEvaluatorModifierEndpoint }
-        PersistHelper.store(headers, as: RequestEvaluatorModifierHeader.storeFileName)
-        PersistHelper.store(endpoints, as: RequestEvaluatorModifierEndpoint.storeFileName)
+public typealias RequestEvaluatorModifier = Modifier & RequestEvaluator & RequestModifier & RequestModifierStorage & Codable
+
+public typealias RequestEvaluatorActionModifier = Modifier & RequestEvaluator & RequestActionModifier & RequestModifierStorage & Codable
+
+extension RequestModifierStorage where Self: Modifier {
+    static func store(_ array: [Codable]) {
+        PersistHelper.store(array.compactMap({ $0 as? Self }), as: Self.storeFileName)
+    }
+
+    static func retrieveFromDisk() -> [Self] {
+        PersistHelper.retrieve(Self.storeFileName, as: [Self].self) ?? []
     }
+}
 
-    func retrieveFromDisk() -> [RequestEvaluatorModifier] {
-        var modifiers = [RequestEvaluatorModifier]()
-        modifiers.append(contentsOf: PersistHelper.retrieve(RequestEvaluatorModifierHeader.storeFileName, as: [RequestEvaluatorModifierHeader].self) ?? [])
-        modifiers.append(contentsOf: PersistHelper.retrieve(RequestEvaluatorModifierEndpoint.storeFileName, as: [RequestEvaluatorModifierEndpoint].self) ?? [])
+extension Array where Element == Modifier {
+    func store() {
+        RequestEvaluatorModifierHeader.store(self)
+        RequestEvaluatorModifierEndpoint.store(self)
+        RequestEvaluatorModifierResponse.store(self)
+    }
+
+    func retrieveFromDisk() -> [Modifier] {
+        var modifiers = [Modifier]()
+        modifiers.append(contentsOf: RequestEvaluatorModifierHeader.retrieveFromDisk())
+        modifiers.append(contentsOf: RequestEvaluatorModifierEndpoint.retrieveFromDisk())
+        modifiers.append(contentsOf: RequestEvaluatorModifierResponse.retrieveFromDisk())
         return modifiers
     }
 }

Sources/NetShears/RequestModifier/RequestEvaluatorModifierEndpoint.swift

@@ -7,7 +7,7 @@
 
 import Foundation
 
-public struct RequestEvaluatorModifierEndpoint: RequestEvaluatorModifier, Equatable, Codable {
+public struct RequestEvaluatorModifierEndpoint: RequestEvaluatorModifier, Equatable {
 
     public var redirectedRequest: RedirectedRequestModel
 

Sources/NetShears/RequestModifier/RequestEvaluatorModifierHeader.swift

@@ -7,7 +7,7 @@
 
 import Foundation
 
-public struct RequestEvaluatorModifierHeader: RequestEvaluatorModifier, Equatable, Codable {
+public struct RequestEvaluatorModifierHeader: RequestEvaluatorModifier, Equatable {
 
     public var header: HeaderModifyModel
 

Sources/NetShears/RequestModifier/RequestEvaluatorModifierResponse.swift

@@ -0,0 +1,31 @@
+//
+//  RequestEvaluatorModifierResponse.swift
+//  
+//
+//  Created by Pouya on 11/4/1400 AP.
+//
+
+import Foundation
+
+public struct RequestEvaluatorModifierResponse: RequestEvaluatorActionModifier, Equatable {
+    public var response: HTTPResponseModifyModel
+
+    public init(response: HTTPResponseModifyModel) {
+        self.response = response
+    }
+
+    public static var storeFileName: String {
+        "Response.txt"
+    }
+
+    public func isActionAllowed(urlRequest: URLRequest) -> Bool {
+        URL(string: response.url) == urlRequest.url && urlRequest.httpMethod?.lowercased() == response.httpMethod.lowercased()
+    }
+
+    public func modify(client: URLProtocolClient?, urlProtocol: URLProtocol) {
+        guard let urlResponse = response.response else { return }
+        client?.urlProtocol(urlProtocol, didLoad: response.data)
+        client?.urlProtocol(urlProtocol, didReceive: urlResponse, cacheStoragePolicy: .notAllowed)
+        client?.urlProtocolDidFinishLoading(urlProtocol)
+    }
+}

Sources/NetShears/URLProtocol/NetworkInterceptorUrlProtocol.swift

@@ -29,7 +29,7 @@ class NetworkInterceptorUrlProtocol: URLProtocol {
         guard NetworkInterceptor.shared.shouldRequestModify(urlRequest: request) else { return false }
 
         if NetworkInterceptorUrlProtocol.property(forKey: Constants.RequestHandledKey, in: request) != nil {
-            return false
+            return actionModifier(forRequest: request) != nil
         }
         return true
     }
@@ -41,8 +41,14 @@ class NetworkInterceptorUrlProtocol: URLProtocol {
     }
 
     override func startLoading() {
+        if let actionModifier = Self.actionModifier(forRequest: request) {
+            actionModifier.modify(client: client, urlProtocol: self)
+            return
+        }
+
         var newRequest = request
-        for modifier in NetShears.shared.config.modifiers where modifier.isActionAllowed(urlRequest: request) {
+        let modifiers = NetShears.shared.config.modifiers.compactMap({ $0 as? RequestEvaluatorModifier })
+        for modifier in modifiers where modifier.isActionAllowed(urlRequest: request) {
             modifier.modify(request: &newRequest)
         }
 
@@ -60,6 +66,13 @@ class NetworkInterceptorUrlProtocol: URLProtocol {
         session = nil
         sessionTask = nil
     }
+
+    class func actionModifier(forRequest request: URLRequest) -> RequestEvaluatorActionModifier? {
+        NetShears.shared.config.modifiers
+            .compactMap({ $0 as? RequestEvaluatorActionModifier })
+            .filter({ $0.isActionAllowed(urlRequest: request) })
+            .last
+    }
 }
 
 extension NetworkInterceptorUrlProtocol: URLSessionDataDelegate {

Tests/NetShearsTests/NetShearsTests.swift

@@ -2,9 +2,80 @@ import XCTest
 @testable import NetShears
 
 final class NetShearsTests: XCTestCase {
-    func testExample() {
-        // This is an example of a functional test case.
-        // Use XCTAssert and related functions to verify your tests produce the correct
-        // results.
+    let expectation = XCTestExpectation()
+
+    override func setUpWithError() throws {
+        try super.setUpWithError()
+        NetShears.shared.startInterceptor()
+    }
+
+    override func tearDownWithError() throws {
+        PersistHelper.clear()
+        NetShears.shared.stopInterceptor()
+        Storage.shared.clearRequests()
+        try super.tearDownWithError()
+    }
+
+    func testModifyResponse() {
+        let modifier = RequestEvaluatorModifierResponse(response: .init(
+            url: "https://example.com/",
+            data: #"{"message": "ok"}"#.data(using: .utf8)!,
+            httpMethod: "POST",
+            headers: ["result": "true"]
+        ))
+
+        NetShears.shared.modify(modifier: modifier)
+
+        var request = URLRequest(url: URL(string: "https://example.com/")!)
+        request.httpMethod = "POST"
+
+        URLSession.shared.dataTask(with: request) { data, response, _ in
+            let dict = try! JSONSerialization.jsonObject(with: data!, options: .fragmentsAllowed) as! [String: String]
+            XCTAssertEqual((response as! HTTPURLResponse).statusCode, 200)
+            XCTAssertEqual((response as! HTTPURLResponse).allHeaderFields["result"] as! String, "true")
+            XCTAssertEqual(dict["message"], "ok")
+            self.expectation.fulfill()
+        }.resume()
+
+        wait(for: [expectation], timeout: 1)
+    }
+
+    func testModifyRequestBeforeModifiyingResponse() {
+        let modifier = RequestEvaluatorModifierResponse(response: .init(
+            url: "https://example.com/sandbox/ok",
+            data: #"{"message": "ok"}"#.data(using: .utf8)!
+        ))
+
+        NetShears.shared.modify(modifier: RequestEvaluatorModifierEndpoint(redirectedRequest: .init(originalUrl: "v1", redirectUrl: "sandbox")))
+        NetShears.shared.modify(modifier: modifier)
+
+        let task = URLSession.shared.dataTask(with: URL(string: "https://example.com/v1/ok")!) { data, response, _ in
+            let dict = try! JSONSerialization.jsonObject(with: data!, options: .fragmentsAllowed) as! [String: String]
+            XCTAssertEqual((response as! HTTPURLResponse).statusCode, 200)
+            XCTAssertEqual(dict["message"], "ok")
+            self.expectation.fulfill()
+        }
+
+        task.resume()
+
+        wait(for: [expectation], timeout: 1)
+    }
+
+    func testMonitorModifiedResponseRequest() {
+        NetShears.shared.startLogger()
+
+        let modifier = RequestEvaluatorModifierResponse(response: .init(
+            url: "https://example.com/",
+            data: #"{"message": "ok"}"#.data(using: .utf8)!
+        ))
+
+        NetShears.shared.modify(modifier: modifier)
+
+        URLSession.shared.dataTask(with: URL(string: "https://example.com/")!) { data, response, _ in
+            XCTAssertTrue(Storage.shared.requests.contains(where: { $0.url == modifier.response.url }))
+            self.expectation.fulfill()
+        }.resume()
+
+        wait(for: [expectation], timeout: 1)
     }
 }