1 #include <iostream>
2 #include <memory>
3 #include <string>
4
5 #include <grpc++/grpc++.h>
6
7 #include "fastnetmon.grpc.pb.h"
8
9 using grpc::Channel;
10 using grpc::ClientContext;
11 using grpc::Status;
12 using fastmitigation::BanListRequest;
13 using fastmitigation::BanListReply;
14 using fastmitigation::Fastnetmon;
15
16 unsigned int client_connection_timeout = 5;
17
18 class FastnetmonClient {
19 public:
FastnetmonClient(std::shared_ptr<Channel> channel)20 FastnetmonClient(std::shared_ptr<Channel> channel) : stub_(Fastnetmon::NewStub(channel)) {}
21
ExecuteBan(std::string host,bool is_ban)22 void ExecuteBan(std::string host, bool is_ban) {
23 ClientContext context;
24 fastmitigation::ExecuteBanRequest request;
25 fastmitigation::ExecuteBanReply reply;
26
27 request.set_ip_address(host);
28
29 Status status;
30
31 if (is_ban) {
32 status = stub_->ExecuteBan(&context, request, &reply);
33 } else {
34 status = stub_->ExecuteUnBan(&context, request, &reply);
35 }
36
37 if (status.ok()) {
38
39 } else {
40 if (status.error_code() == grpc::DEADLINE_EXCEEDED) {
41 std::cerr << "Could not connect to API server. Timeout exceed" << std::endl;
42 return;
43 } else {
44 std::cerr << "RPC failed " + status.error_message() << std::endl;
45 return;
46 }
47 }
48 }
49
GetBanList()50 void GetBanList() {
51 // This request haven't any useful data
52 BanListRequest request;
53
54 // Container for the data we expect from the server.
55 BanListReply reply;
56
57 // Context for the client. It could be used to convey extra information to
58 // the server and/or tweak certain RPC behaviors.
59 ClientContext context;
60
61 // Set timeout for API
62 std::chrono::system_clock::time_point deadline =
63 std::chrono::system_clock::now() + std::chrono::seconds(client_connection_timeout);
64
65 context.set_deadline(deadline);
66
67 // The actual RPC.
68 auto announces_list = stub_->GetBanlist(&context, request);
69
70 while (announces_list->Read(&reply)) {
71 std::cout << reply.ip_address() << std::endl;
72 }
73
74 // Get status and handle errors
75 auto status = announces_list->Finish();
76
77 if (!status.ok()) {
78 if (status.error_code() == grpc::DEADLINE_EXCEEDED) {
79 std::cerr << "Could not connect to API server. Timeout exceed" << std::endl;
80 return;
81 } else {
82 std::cerr << "RPC failed " + status.error_message();
83 return;
84 }
85 }
86 }
87
88 private:
89 std::unique_ptr<Fastnetmon::Stub> stub_;
90 };
91
silent_logging_function(gpr_log_func_args * args)92 void silent_logging_function(gpr_log_func_args *args) {
93 // We do not want any logging here
94 }
95
main(int argc,char ** argv)96 int main(int argc, char** argv) {
97 if (argc <= 1) {
98 std::cerr << "Please provide request" << std::endl;
99 return 1;
100 }
101
102 gpr_set_log_function(silent_logging_function);
103
104 // Instantiate the client. It requires a channel, out of which the actual RPCs
105 // are created. This channel models a connection to an endpoint (in this case,
106 // localhost at port 50051). We indicate that the channel isn't authenticated
107 // (use of InsecureCredentials()).
108 FastnetmonClient fastnetmon( grpc::CreateChannel("localhost:50052", grpc::InsecureCredentials()));
109
110 std::string request_command = argv[1];
111
112 if (request_command == "get_banlist") {
113 fastnetmon.GetBanList();
114 } else if (request_command == "ban" or request_command == "unban") {
115 if (argc < 2) {
116 std::cerr << "Please provide IP for action" << std::endl;
117 return(1);
118 }
119
120 std::string ip_for_ban = argv[2];
121
122 if (request_command == "ban") {
123 fastnetmon.ExecuteBan(ip_for_ban, true);
124 } else {
125 fastnetmon.ExecuteBan(ip_for_ban, false);
126 }
127 } else {
128 std::cerr << "Unknown command" << std::endl;
129 }
130
131 return 0;
132 }
133