]> git.stg.codes - stg.git/blob - projects/stargazer/traffcounter_impl.h
78027cb57d3d7d5eb8b08448b3b4e81843c9d65a
[stg.git] / projects / stargazer / traffcounter_impl.h
1 /*
2  *    This program is free software; you can redistribute it and/or modify
3  *    it under the terms of the GNU General Public License as published by
4  *    the Free Software Foundation; either version 2 of the License, or
5  *    (at your option) any later version.
6  *
7  *    This program is distributed in the hope that it will be useful,
8  *    but WITHOUT ANY WARRANTY; without even the implied warranty of
9  *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
10  *    GNU General Public License for more details.
11  *
12  *    You should have received a copy of the GNU General Public License
13  *    along with this program; if not, write to the Free Software
14  *    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
15  */
16
17 /*
18  *    Author : Boris Mikhailenko <stg34@stargazer.dp.ua>
19  */
20
21 #pragma once
22
23 #include "stg/traffcounter.h"
24 #include "stg/logger.h"
25 #include "stg/raw_ip_packet.h"
26 #include "stg/noncopyable.h"
27 #include "stg/notifer.h"
28 #include "actions.h"
29 #include "eventloop.h"
30 #include "user_impl.h"
31
32 #include <ctime>
33 #include <list>
34 #include <map>
35 #include <string>
36 #include <cstdint>
37
38 #include <pthread.h>
39
40 #define PROTOMAX    (5)
41
42 namespace STG
43 {
44
45 class UsersImpl;
46
47 //-----------------------------------------------------------------------------
48 struct Rule {
49 uint32_t    ip;             // IP
50 uint32_t    mask;           // Network mask
51 uint16_t    port1;          // Min port
52 uint16_t    port2;          // Max port
53 uint8_t     proto;          // Protocol
54 uint32_t    dir;            // Direction
55 };
56 //-----------------------------------------------------------------------------
57 struct PacketExtraData {
58     PacketExtraData()
59         : flushTime(0),
60           updateTime(0),
61           userU(NULL),
62           userUPresent(false),
63           userD(NULL),
64           userDPresent(false),
65           dirU(DIR_NUM),
66           dirD(DIR_NUM),
67           lenU(0),
68           lenD(0)
69     {}
70
71     time_t      flushTime;          // Last flush time
72     time_t      updateTime;         // Last update time
73     UserImpl * userU;              // Uploader
74     bool        userUPresent;       // Uploader is registered user
75     UserImpl * userD;              // Downloader
76     bool        userDPresent;       // Downloader is registered user
77     int         dirU;               // Upload direction
78     int         dirD;               // Download direction
79     uint32_t    lenU;               // Upload length
80     uint32_t    lenD;               // Download length
81 };
82 //-----------------------------------------------------------------------------
83 class TraffCounterImpl;
84 //-----------------------------------------------------------------------------
85 class TRF_IP_BEFORE: public PropertyNotifierBase<uint32_t> {
86 public:
87                 TRF_IP_BEFORE(TraffCounterImpl & t, UserImpl * u)
88                     : PropertyNotifierBase<uint32_t>(),
89                       traffCnt(t),
90                       user(u)
91                 {}
92                 TRF_IP_BEFORE(const TRF_IP_BEFORE & rvalue)
93                     : PropertyNotifierBase<uint32_t>(),
94                       traffCnt(rvalue.traffCnt),
95                       user(rvalue.user)
96                 {}
97     void        Notify(const uint32_t & oldValue, const uint32_t & newValue);
98     void        SetUser(UserImpl * u) { user = u; }
99     UserImpl * GetUser() const { return user; }
100
101 private:
102     TRF_IP_BEFORE & operator=(const TRF_IP_BEFORE & rvalue);
103
104     TraffCounterImpl & traffCnt;
105     UserImpl * user;
106 };
107 //-----------------------------------------------------------------------------
108 class TRF_IP_AFTER: public PropertyNotifierBase<uint32_t> {
109 public:
110                 TRF_IP_AFTER(TraffCounterImpl & t, UserImpl * u)
111                     : PropertyNotifierBase<uint32_t>(),
112                       traffCnt(t),
113                       user(u)
114                 {}
115                 TRF_IP_AFTER(const TRF_IP_AFTER & rvalue)
116                     : PropertyNotifierBase<uint32_t>(),
117                       traffCnt(rvalue.traffCnt),
118                       user(rvalue.user)
119                 {}
120     void        Notify(const uint32_t & oldValue, const uint32_t & newValue);
121     void        SetUser(UserImpl * u) { user = u; }
122     UserImpl * GetUser() const { return user; }
123 private:
124     TRF_IP_AFTER & operator=(const TRF_IP_AFTER & rvalue);
125
126     TraffCounterImpl & traffCnt;
127     UserImpl * user;
128 };
129
130 using UserImplPtr = UserImpl*;
131 //-----------------------------------------------------------------------------
132 class ADD_USER_NONIFIER: public NotifierBase<UserImplPtr> {
133 public:
134             explicit ADD_USER_NONIFIER(TraffCounterImpl & t) :
135                 NotifierBase<UserImplPtr>(),
136                 traffCnt(t)
137             {}
138     virtual ~ADD_USER_NONIFIER() {}
139     void    Notify(const UserImplPtr & user);
140
141 private:
142     ADD_USER_NONIFIER(const ADD_USER_NONIFIER & rvalue);
143     ADD_USER_NONIFIER & operator=(const ADD_USER_NONIFIER & rvalue);
144
145     TraffCounterImpl & traffCnt;
146 };
147 //-----------------------------------------------------------------------------
148 class DEL_USER_NONIFIER: public NotifierBase<UserImplPtr> {
149 public:
150             explicit DEL_USER_NONIFIER(TraffCounterImpl & t) :
151                 NotifierBase<UserImplPtr>(),
152                 traffCnt(t)
153             {}
154     virtual ~DEL_USER_NONIFIER() {}
155     void    Notify(const UserImplPtr & user);
156
157 private:
158     DEL_USER_NONIFIER(const DEL_USER_NONIFIER & rvalue);
159     DEL_USER_NONIFIER & operator=(const DEL_USER_NONIFIER & rvalue);
160
161     TraffCounterImpl & traffCnt;
162 };
163 //-----------------------------------------------------------------------------
164 class TraffCounterImpl : public TraffCounter {
165     friend class ADD_USER_NONIFIER;
166     friend class DEL_USER_NONIFIER;
167     friend class TRF_IP_BEFORE;
168     friend class TRF_IP_AFTER;
169     public:
170         TraffCounterImpl(UsersImpl * users, const std::string & rulesFileName);
171         ~TraffCounterImpl();
172
173         int         Reload();
174         int         Start();
175         int         Stop();
176
177         void        process(const RawPacket & rawPacket) override;
178         void        SetMonitorDir(const std::string & monitorDir);
179
180         size_t      rulesCount() const override { return rules.size(); }
181
182     private:
183         bool        ParseAddress(const char * ta, Rule * rule) const;
184         uint32_t    CalcMask(uint32_t msk) const;
185         void        FreeRules();
186         bool        ReadRules(bool test = false);
187
188         static void * Run(void * data);
189
190         void        DeterminateDir(const RawPacket & packet,
191                                    int * dirU, // Direction for upload
192                                    int * dirD) const; // Direction for download
193
194         void        FlushAndRemove();
195
196         void        AddUser(UserImpl * user);
197         void        DelUser(uint32_t uip);
198         void        SetUserNotifiers(UserImpl * user);
199         void        UnSetUserNotifiers(UserImpl * user);
200
201         typedef std::list<Rule>::iterator rule_iter;
202
203         std::list<Rule>          rules;
204
205         typedef std::map<RawPacket, PacketExtraData> Packets;
206         typedef Packets::iterator pp_iter;
207         typedef std::multimap<uint32_t, pp_iter> Index;
208         typedef Index::iterator ip2p_iter;
209         typedef Index::const_iterator ip2p_citer;
210
211         Packets packets; // Packets tree
212
213         Index ip2packets; // IP-to-Packet index
214
215         std::string              dirName[DIR_NUM + 1];
216
217         Logger &             WriteServLog;
218         std::string              rulesFileName;
219
220         std::string              monitorDir;
221         bool                     monitoring;
222         time_t                   touchTimeP;
223
224         UsersImpl *             users;
225
226         bool                     running;
227         bool                     stopped;
228         pthread_mutex_t          mutex;
229         pthread_t                thread;
230
231         std::list<TRF_IP_BEFORE> ipBeforeNotifiers;
232         std::list<TRF_IP_AFTER>  ipAfterNotifiers;
233
234         ADD_USER_NONIFIER        addUserNotifier;
235         DEL_USER_NONIFIER        delUserNotifier;
236 };
237 //-----------------------------------------------------------------------------
238 inline
239 void TRF_IP_BEFORE::Notify(const uint32_t & oldValue, const uint32_t &)
240 {
241 // User changes his address. Remove old IP
242 if (!oldValue)
243     return;
244
245 EVENT_LOOP_SINGLETON::GetInstance().Enqueue(traffCnt, &TraffCounterImpl::DelUser, oldValue);
246 }
247 //-----------------------------------------------------------------------------
248 inline
249 void TRF_IP_AFTER::Notify(const uint32_t &, const uint32_t & newValue)
250 {
251 // User changes his address. Add new IP
252 if (!newValue)
253     return;
254
255 EVENT_LOOP_SINGLETON::GetInstance().Enqueue(traffCnt, &TraffCounterImpl::AddUser, user);
256 }
257 //-----------------------------------------------------------------------------
258 inline
259 void ADD_USER_NONIFIER::Notify(const UserImplPtr & user)
260 {
261 EVENT_LOOP_SINGLETON::GetInstance().Enqueue(traffCnt, &TraffCounterImpl::SetUserNotifiers, user);
262 }
263 //-----------------------------------------------------------------------------
264 inline
265 void DEL_USER_NONIFIER::Notify(const UserImplPtr & user)
266 {
267 EVENT_LOOP_SINGLETON::GetInstance().Enqueue(traffCnt, &TraffCounterImpl::UnSetUserNotifiers, user);
268 EVENT_LOOP_SINGLETON::GetInstance().Enqueue(traffCnt, &TraffCounterImpl::DelUser, user->GetCurrIP());
269 }
270 //-----------------------------------------------------------------------------
271 }