From 1f46e83507a259bb92aad50bcd2d4ba21a27efa9 Mon Sep 17 00:00:00 2001
From: Pavel Odintsov <pavel.odintsov@gmail.com>
Date: Thu, 15 Oct 2015 22:27:51 +0200
Subject: [PATCH] Add unban action

---
 src/api_client.cpp   | 23 +++++++++++++++++------
 src/api_server.cpp   |  6 ++++++
 src/fastnetmon.proto |  1 +
 3 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/api_client.cpp b/src/api_client.cpp
index b2b24326..c5af15ff 100644
--- a/src/api_client.cpp
+++ b/src/api_client.cpp
@@ -19,15 +19,21 @@ class GreeterClient {
     public:
         GreeterClient(std::shared_ptr<Channel> channel) : stub_(Fastnetmon::NewStub(channel)) {}
 
-        void ExecuteBan(std::string host) {
+        void ExecuteBan(std::string host, bool is_ban) {
             ClientContext context;
             fastmitigation::ExecuteBanRequest request;
             fastmitigation::ExecuteBanReply reply;
 
             request.set_ip_address(host);
 
-            Status status = stub_->ExecuteBan(&context, request, &reply);
-
+            Status status;
+    
+            if (is_ban) {
+                status = stub_->ExecuteBan(&context, request, &reply);
+            } else {
+                status = stub_->ExecuteUnBan(&context, request, &reply);
+            }
+    
             if (status.ok()) {
                 
             } else {
@@ -105,14 +111,19 @@ int main(int argc, char** argv) {
 
     if (request_command == "get_banlist") {
         greeter.GetBanList();
-    } else if (request_command == "ban") {
+    } else if (request_command == "ban" or request_command == "unban") {
         if (argc < 2) {
-            std::cerr << "Please provide banned IP" << std::endl;
+            std::cerr << "Please provide IP for action" << std::endl;
             return(1);
         }
 
         std::string ip_for_ban = argv[2];
-        greeter.ExecuteBan(ip_for_ban); 
+
+        if (request_command == "ban") {
+            greeter.ExecuteBan(ip_for_ban, true);
+        } else {
+            greeter.ExecuteBan(ip_for_ban, false);
+        }
     } else {
         std::cerr << "Unknown command" << std::endl;
     }
diff --git a/src/api_server.cpp b/src/api_server.cpp
index b3b54361..63f74dea 100644
--- a/src/api_server.cpp
+++ b/src/api_server.cpp
@@ -35,6 +35,12 @@ class GreeterServiceImpl final : public Fastnetmon::Service {
 
         return Status::OK;
     }
+
+    Status ExecuteUnBan(ServerContext* context, const fastmitigation::ExecuteBanRequest* request, fastmitigation::ExecuteBanReply* reply) override {
+        std::cout << "We asked for unban for IP: " << request->ip_address() << std::endl;
+
+        return Status::OK;
+    }
 };
 
 void RunServer() {
diff --git a/src/fastnetmon.proto b/src/fastnetmon.proto
index 545d0a87..bc88181e 100644
--- a/src/fastnetmon.proto
+++ b/src/fastnetmon.proto
@@ -5,6 +5,7 @@ package fastmitigation;
 service Fastnetmon {
     rpc GetBanlist(BanListRequest) returns (stream BanListReply) {}
     rpc ExecuteBan(ExecuteBanRequest) returns (ExecuteBanReply) {}
+    rpc ExecuteUnBan(ExecuteBanRequest) returns (ExecuteBanReply) {}
 }
 
 // We could not create RPC method without params