gitscope
/
← marzneshin / app/db/crud.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/db/crud.py · v0.6.4
Source
1 12b255b Dawsh 2y ago import json
2 34a7612 Dawsh 2y ago import secrets
3 e6e6467 Dawsh 2y ago from collections import defaultdict
4 e6e6467 Dawsh 2y ago from datetime import datetime, timedelta, timezone
5 a47eb4c Dawsh 2y ago from enum import Enum
6 4708fcb Dawsh 2y ago from types import NoneType
7 2de9d72 Dawsh 2y ago from typing import List, Optional, Tuple, Union
8 a47eb4c Dawsh 2y ago
9 9ca312d Dawsh 2y ago from sqlalchemy import and_, delete, update, select, func, cast, Date
10 a47eb4c Dawsh 2y ago from sqlalchemy.orm import Session
11 a47eb4c Dawsh 2y ago
12 01cdf2e Dawsh 2y ago from app.db.models import (
13 01cdf2e Dawsh 2y ago JWT,
14 01cdf2e Dawsh 2y ago TLS,
15 01cdf2e Dawsh 2y ago Admin,
16 01cdf2e Dawsh 2y ago Node,
17 01cdf2e Dawsh 2y ago NodeUserUsage,
18 01cdf2e Dawsh 2y ago InboundHost,
19 01cdf2e Dawsh 2y ago Service,
20 01cdf2e Dawsh 2y ago Inbound,
21 01cdf2e Dawsh 2y ago System,
22 01cdf2e Dawsh 2y ago User,
23 f5c168d Mardin 2y ago Backend,
24 01cdf2e Dawsh 2y ago )
25 4708fcb Dawsh 2y ago from app.models.admin import AdminCreate, AdminPartialModify
26 f1ab16c Dawsh 2y ago from app.models.node import (
27 f1ab16c Dawsh 2y ago NodeCreate,
28 f1ab16c Dawsh 2y ago NodeModify,
29 f1ab16c Dawsh 2y ago NodeStatus,
30 f1ab16c Dawsh 2y ago )
31 a47eb4c Dawsh 2y ago from app.models.proxy import InboundHost as InboundHostModify
32 2de9d72 Dawsh 2y ago from app.models.service import Service as ServiceModify, ServiceCreate
33 6a4fee1 Dawsh 2y ago from app.models.system import TrafficUsageSeries
34 01cdf2e Dawsh 2y ago from app.models.user import (
35 01cdf2e Dawsh 2y ago UserCreate,
36 d377c4c Dawsh 2y ago UserDataUsageResetStrategy,
37 01cdf2e Dawsh 2y ago UserModify,
38 01cdf2e Dawsh 2y ago UserStatus,
39 d377c4c Dawsh 2y ago UserExpireStrategy,
40 e6e6467 Dawsh 2y ago UserNodeUsageSeries,
41 e6e6467 Dawsh 2y ago UserUsageSeriesResponse,
42 f1ab16c Dawsh 2y ago )
43 a47eb4c Dawsh 2y ago
44 a47eb4c Dawsh 2y ago
45 a47eb4c Dawsh 2y ago def add_default_hosts(db: Session, inbounds: List[Inbound]):
46 01cdf2e Dawsh 2y ago hosts = [
47 01cdf2e Dawsh 2y ago InboundHost(
48 01cdf2e Dawsh 2y ago remark="🚀 Marz ({USERNAME}) [{PROTOCOL} - {TRANSPORT}]",
49 01cdf2e Dawsh 2y ago address="{SERVER_IP}",
50 01cdf2e Dawsh 2y ago inbound=i,
51 01cdf2e Dawsh 2y ago )
52 01cdf2e Dawsh 2y ago for i in inbounds
53 01cdf2e Dawsh 2y ago ]
54 a47eb4c Dawsh 2y ago db.add_all(hosts)
55 a47eb4c Dawsh 2y ago db.commit()
56 a47eb4c Dawsh 2y ago
57 a47eb4c Dawsh 2y ago
58 f5c168d Mardin 2y ago def ensure_node_backends(db: Session, backends, node_id: int):
59 f5c168d Mardin 2y ago old_backends = db.query(Backend).where(Backend.node_id == node_id)
60 f5c168d Mardin 2y ago for backend in old_backends:
61 f5c168d Mardin 2y ago db.delete(backend)
62 f5c168d Mardin 2y ago backends = [
63 f5c168d Mardin 2y ago Backend(
64 f5c168d Mardin 2y ago name=backend.name,
65 f5c168d Mardin 2y ago backend_type=backend.type,
66 f5c168d Mardin 2y ago version=backend.version,
67 f5c168d Mardin 2y ago node_id=node_id,
68 f5c168d Mardin 2y ago )
69 f5c168d Mardin 2y ago for backend in backends
70 f5c168d Mardin 2y ago ]
71 f5c168d Mardin 2y ago db.add_all(backends)
72 f5c168d Mardin 2y ago db.flush()
73 f5c168d Mardin 2y ago
74 f5c168d Mardin 2y ago
75 f5c168d Mardin 2y ago def ensure_node_inbounds(db: Session, inbounds: List[Inbound], node_id: int):
76 01cdf2e Dawsh 2y ago current_tags = [
77 01cdf2e Dawsh 2y ago i[0]
78 01cdf2e Dawsh 2y ago for i in db.execute(
79 01cdf2e Dawsh 2y ago select(Inbound.tag).filter(Inbound.node_id == node_id)
80 01cdf2e Dawsh 2y ago ).all()
81 01cdf2e Dawsh 2y ago ]
82 512d183 Mrdn 2y ago updated_tags = set(i.tag for i in list(inbounds))
83 512d183 Mrdn 2y ago inbound_additions, tag_deletions = list(), set()
84 512d183 Mrdn 2y ago for tag in current_tags:
85 512d183 Mrdn 2y ago if tag not in updated_tags:
86 512d183 Mrdn 2y ago tag_deletions.add(tag)
87 f5c168d Mardin 2y ago removals = db.query(Inbound).where(
88 01cdf2e Dawsh 2y ago and_(Inbound.node_id == node_id, Inbound.tag.in_(tag_deletions))
89 01cdf2e Dawsh 2y ago )
90 f5c168d Mardin 2y ago for i in removals:
91 512d183 Mrdn 2y ago db.delete(i)
92 512d183 Mrdn 2y ago
93 512d183 Mrdn 2y ago for inb in inbounds:
94 512d183 Mrdn 2y ago if inb.tag in current_tags:
95 01cdf2e Dawsh 2y ago stmt = (
96 01cdf2e Dawsh 2y ago update(Inbound)
97 f1ab16c Dawsh 2y ago .where(
98 f1ab16c Dawsh 2y ago and_(Inbound.node_id == node_id, Inbound.tag == inb.tag)
99 f1ab16c Dawsh 2y ago )
100 f1ab16c Dawsh 2y ago .values(
101 f1ab16c Dawsh 2y ago protocol=json.loads(inb.config)["protocol"],
102 f1ab16c Dawsh 2y ago config=inb.config,
103 f1ab16c Dawsh 2y ago )
104 a47eb4c Dawsh 2y ago )
105 512d183 Mrdn 2y ago db.execute(stmt)
106 512d183 Mrdn 2y ago else:
107 512d183 Mrdn 2y ago inbound_additions.append(inb)
108 01cdf2e Dawsh 2y ago new_inbounds = [
109 01cdf2e Dawsh 2y ago Inbound(
110 01cdf2e Dawsh 2y ago tag=inb.tag,
111 01cdf2e Dawsh 2y ago protocol=json.loads(inb.config)["protocol"],
112 01cdf2e Dawsh 2y ago config=inb.config,
113 01cdf2e Dawsh 2y ago node_id=node_id,
114 01cdf2e Dawsh 2y ago )
115 01cdf2e Dawsh 2y ago for inb in inbound_additions
116 01cdf2e Dawsh 2y ago ]
117 a47eb4c Dawsh 2y ago db.add_all(new_inbounds)
118 a47eb4c Dawsh 2y ago db.flush()
119 a47eb4c Dawsh 2y ago for i in new_inbounds:
120 2de9d72 Dawsh 2y ago db.refresh(i)
121 a47eb4c Dawsh 2y ago add_default_hosts(db, new_inbounds)
122 a47eb4c Dawsh 2y ago db.commit()
123 a47eb4c Dawsh 2y ago
124 a47eb4c Dawsh 2y ago
125 01cdf2e Dawsh 2y ago def get_node_users(
126 f1ab16c Dawsh 2y ago db: Session,
127 d377c4c Dawsh 2y ago node_id: int,
128 01cdf2e Dawsh 2y ago ):
129 01cdf2e Dawsh 2y ago query = (
130 01cdf2e Dawsh 2y ago db.query(User.id, User.username, User.key, Inbound)
131 01cdf2e Dawsh 2y ago .distinct()
132 01cdf2e Dawsh 2y ago .join(Inbound.services)
133 01cdf2e Dawsh 2y ago .join(Service.users)
134 01cdf2e Dawsh 2y ago .filter(Inbound.node_id == node_id)
135 d377c4c Dawsh 2y ago .filter(User.activated == True)
136 01cdf2e Dawsh 2y ago )
137 a47eb4c Dawsh 2y ago return query.all()
138 a47eb4c Dawsh 2y ago
139 12b255b Dawsh 2y ago
140 a47eb4c Dawsh 2y ago def get_user_hosts(db: Session, user_id: int):
141 01cdf2e Dawsh 2y ago return (
142 01cdf2e Dawsh 2y ago db.query(InboundHost)
143 01cdf2e Dawsh 2y ago .distinct()
144 01cdf2e Dawsh 2y ago .join(User.services)
145 01cdf2e Dawsh 2y ago .join(Service.inbounds)
146 01cdf2e Dawsh 2y ago .join(Inbound.hosts)
147 01cdf2e Dawsh 2y ago .filter(User.id == user_id)
148 01cdf2e Dawsh 2y ago .all()
149 01cdf2e Dawsh 2y ago )
150 a47eb4c Dawsh 2y ago
151 12b255b Dawsh 2y ago
152 a47eb4c Dawsh 2y ago def get_inbound_hosts(db: Session, inbound_id: int) -> List[InboundHost]:
153 f1ab16c Dawsh 2y ago return (
154 f1ab16c Dawsh 2y ago db.query(InboundHost)
155 f1ab16c Dawsh 2y ago .filter(InboundHost.inbound_id == inbound_id)
156 f1ab16c Dawsh 2y ago .all()
157 f1ab16c Dawsh 2y ago )
158 a47eb4c Dawsh 2y ago
159 12b255b Dawsh 2y ago
160 a47eb4c Dawsh 2y ago def get_all_inbounds(db: Session):
161 a47eb4c Dawsh 2y ago return db.query(Inbound).all()
162 a47eb4c Dawsh 2y ago
163 a47eb4c Dawsh 2y ago
164 2de9d72 Dawsh 2y ago def get_inbound(db: Session, inbound_id: int) -> Inbound | None:
165 2de9d72 Dawsh 2y ago return db.query(Inbound).filter(Inbound.id == inbound_id).first()
166 a47eb4c Dawsh 2y ago
167 a47eb4c Dawsh 2y ago
168 2de9d72 Dawsh 2y ago def get_host(db: Session, host_id) -> InboundHost:
169 2de9d72 Dawsh 2y ago return db.query(InboundHost).filter(InboundHost.id == host_id).first()
170 a47eb4c Dawsh 2y ago
171 a47eb4c Dawsh 2y ago
172 a47eb4c Dawsh 2y ago def add_host(db: Session, inbound: Inbound, host: InboundHostModify):
173 2de9d72 Dawsh 2y ago host = InboundHost(
174 2de9d72 Dawsh 2y ago remark=host.remark,
175 2de9d72 Dawsh 2y ago address=host.address,
176 2de9d72 Dawsh 2y ago port=host.port,
177 2de9d72 Dawsh 2y ago path=host.path,
178 2de9d72 Dawsh 2y ago sni=host.sni,
179 2de9d72 Dawsh 2y ago host=host.host,
180 2de9d72 Dawsh 2y ago inbound=inbound,
181 2de9d72 Dawsh 2y ago security=host.security,
182 244c9d8 Dawsh 2y ago alpn=host.alpn.value,
183 01cdf2e Dawsh 2y ago fingerprint=host.fingerprint,
184 0b152b5 Mrdn 2y ago fragment=host.fragment.model_dump() if host.fragment else None,
185 d636f3f Dawsh 2y ago udp_noises=host.udp_noises,
186 d636f3f Dawsh 2y ago http_headers=host.http_headers,
187 d636f3f Dawsh 2y ago mtu=host.mtu,
188 d636f3f Dawsh 2y ago dns_servers=host.dns_servers,
189 b584c6c Dawsh 2y ago allowed_ips=host.allowed_ips,
190 0b152b5 Mrdn 2y ago mux=host.mux,
191 dd32f70 Dawsh 2y ago allowinsecure=host.allowinsecure,
192 e422d72 Dawsh 2y ago weight=host.weight,
193 a47eb4c Dawsh 2y ago )
194 2de9d72 Dawsh 2y ago inbound.hosts.append(host)
195 a47eb4c Dawsh 2y ago db.commit()
196 2de9d72 Dawsh 2y ago db.refresh(host)
197 2de9d72 Dawsh 2y ago return host
198 2de9d72 Dawsh 2y ago
199 2de9d72 Dawsh 2y ago
200 2de9d72 Dawsh 2y ago def update_host(db: Session, db_host: InboundHost, host: InboundHostModify):
201 2de9d72 Dawsh 2y ago db_host.remark = host.remark
202 2de9d72 Dawsh 2y ago db_host.address = host.address
203 2de9d72 Dawsh 2y ago db_host.port = host.port
204 2de9d72 Dawsh 2y ago db_host.path = host.path
205 2de9d72 Dawsh 2y ago db_host.sni = host.sni
206 2de9d72 Dawsh 2y ago db_host.host = host.host
207 2de9d72 Dawsh 2y ago db_host.security = host.security
208 244c9d8 Dawsh 2y ago db_host.alpn = host.alpn.value
209 2de9d72 Dawsh 2y ago db_host.fingerprint = host.fingerprint
210 0b152b5 Mrdn 2y ago db_host.fragment = host.fragment.model_dump() if host.fragment else None
211 0b152b5 Mrdn 2y ago db_host.mux = host.mux
212 a88ed47 Mardin 2y ago db_host.is_disabled = host.is_disabled
213 dd32f70 Dawsh 2y ago db_host.allowinsecure = host.allowinsecure
214 d636f3f Dawsh 2y ago db_host.udp_noises = host.udp_noises
215 d636f3f Dawsh 2y ago db_host.http_headers = host.http_headers
216 d636f3f Dawsh 2y ago db_host.mtu = host.mtu
217 d636f3f Dawsh 2y ago db_host.dns_servers = host.dns_servers
218 b584c6c Dawsh 2y ago db_host.allowed_ips = host.allowed_ips
219 e422d72 Dawsh 2y ago db_host.weight = host.weight
220 2de9d72 Dawsh 2y ago db.commit()
221 2de9d72 Dawsh 2y ago db.refresh(db_host)
222 2de9d72 Dawsh 2y ago return db_host
223 a47eb4c Dawsh 2y ago
224 a47eb4c Dawsh 2y ago
225 a47eb4c Dawsh 2y ago def get_user(db: Session, username: str):
226 a47eb4c Dawsh 2y ago return db.query(User).filter(User.username == username).first()
227 a47eb4c Dawsh 2y ago
228 a47eb4c Dawsh 2y ago
229 a47eb4c Dawsh 2y ago def get_user_by_id(db: Session, user_id: int):
230 a47eb4c Dawsh 2y ago return db.query(User).filter(User.id == user_id).first()
231 a47eb4c Dawsh 2y ago
232 a47eb4c Dawsh 2y ago
233 01cdf2e Dawsh 2y ago UsersSortingOptions = Enum(
234 01cdf2e Dawsh 2y ago "UsersSortingOptions",
235 01cdf2e Dawsh 2y ago {
236 01cdf2e Dawsh 2y ago "username": User.username.asc(),
237 01cdf2e Dawsh 2y ago "used_traffic": User.used_traffic.asc(),
238 01cdf2e Dawsh 2y ago "data_limit": User.data_limit.asc(),
239 d377c4c Dawsh 2y ago "expire": User.expire_date.asc(),
240 01cdf2e Dawsh 2y ago "created_at": User.created_at.asc(),
241 01cdf2e Dawsh 2y ago "-username": User.username.desc(),
242 01cdf2e Dawsh 2y ago "-used_traffic": User.used_traffic.desc(),
243 01cdf2e Dawsh 2y ago "-data_limit": User.data_limit.desc(),
244 d377c4c Dawsh 2y ago "-expire": User.expire_date.desc(),
245 01cdf2e Dawsh 2y ago "-created_at": User.created_at.desc(),
246 01cdf2e Dawsh 2y ago },
247 01cdf2e Dawsh 2y ago )
248 01cdf2e Dawsh 2y ago
249 01cdf2e Dawsh 2y ago
250 01cdf2e Dawsh 2y ago def get_users(
251 01cdf2e Dawsh 2y ago db: Session,
252 01cdf2e Dawsh 2y ago offset: Optional[int] = None,
253 01cdf2e Dawsh 2y ago limit: Optional[int] = None,
254 01cdf2e Dawsh 2y ago usernames: Optional[List[str]] = None,
255 01cdf2e Dawsh 2y ago sort: Optional[List[UsersSortingOptions]] = None,
256 01cdf2e Dawsh 2y ago admin: Optional[Admin] = None,
257 d377c4c Dawsh 2y ago reset_strategy: Optional[Union[UserDataUsageResetStrategy, list]] = None,
258 d377c4c Dawsh 2y ago expire_strategy: (
259 d377c4c Dawsh 2y ago UserExpireStrategy | list[UserExpireStrategy] | None
260 d377c4c Dawsh 2y ago ) = None,
261 d377c4c Dawsh 2y ago is_active: bool | None = None,
262 d377c4c Dawsh 2y ago activated: bool | None = None,
263 d377c4c Dawsh 2y ago expired: bool | None = None,
264 d377c4c Dawsh 2y ago data_limit_reached: bool | None = None,
265 4708fcb Dawsh 2y ago enabled: bool | None = None,
266 01cdf2e Dawsh 2y ago ) -> Union[List[User], Tuple[List[User], int]]:
267 24826a2 Dawsh 2y ago query = db.query(User).filter(User.removed == False)
268 a47eb4c Dawsh 2y ago
269 a47eb4c Dawsh 2y ago if usernames:
270 a47eb4c Dawsh 2y ago if len(usernames) == 1:
271 a47eb4c Dawsh 2y ago query = query.filter(User.username.ilike(f"%{usernames[0]}%"))
272 a47eb4c Dawsh 2y ago else:
273 a47eb4c Dawsh 2y ago query = query.filter(User.username.in_(usernames))
274 a47eb4c Dawsh 2y ago
275 a47eb4c Dawsh 2y ago if reset_strategy:
276 a47eb4c Dawsh 2y ago if isinstance(reset_strategy, list):
277 f1ab16c Dawsh 2y ago query = query.filter(
278 f1ab16c Dawsh 2y ago User.data_limit_reset_strategy.in_(reset_strategy)
279 f1ab16c Dawsh 2y ago )
280 a47eb4c Dawsh 2y ago else:
281 f1ab16c Dawsh 2y ago query = query.filter(
282 f1ab16c Dawsh 2y ago User.data_limit_reset_strategy == reset_strategy
283 f1ab16c Dawsh 2y ago )
284 a47eb4c Dawsh 2y ago
285 d377c4c Dawsh 2y ago if expire_strategy:
286 d377c4c Dawsh 2y ago if isinstance(expire_strategy, list):
287 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy.in_(expire_strategy))
288 d377c4c Dawsh 2y ago else:
289 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy == expire_strategy)
290 d377c4c Dawsh 2y ago
291 d377c4c Dawsh 2y ago if isinstance(is_active, bool):
292 d377c4c Dawsh 2y ago query = query.filter(User.is_active == is_active)
293 d377c4c Dawsh 2y ago
294 d377c4c Dawsh 2y ago if isinstance(activated, bool):
295 d377c4c Dawsh 2y ago query = query.filter(User.activated == activated)
296 d377c4c Dawsh 2y ago
297 d377c4c Dawsh 2y ago if isinstance(expired, bool):
298 d377c4c Dawsh 2y ago query = query.filter(User.expired == expired)
299 d377c4c Dawsh 2y ago
300 d377c4c Dawsh 2y ago if isinstance(data_limit_reached, bool):
301 d377c4c Dawsh 2y ago query = query.filter(User.data_limit_reached == data_limit_reached)
302 d377c4c Dawsh 2y ago
303 4708fcb Dawsh 2y ago if isinstance(enabled, bool):
304 4708fcb Dawsh 2y ago query = query.filter(User.enabled == enabled)
305 4708fcb Dawsh 2y ago
306 a47eb4c Dawsh 2y ago if admin:
307 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
308 a47eb4c Dawsh 2y ago
309 a47eb4c Dawsh 2y ago if sort:
310 a47eb4c Dawsh 2y ago query = query.order_by(*(opt.value for opt in sort))
311 a47eb4c Dawsh 2y ago
312 a47eb4c Dawsh 2y ago if offset:
313 a47eb4c Dawsh 2y ago query = query.offset(offset)
314 d377c4c Dawsh 2y ago
315 a47eb4c Dawsh 2y ago if limit:
316 a47eb4c Dawsh 2y ago query = query.limit(limit)
317 a47eb4c Dawsh 2y ago
318 a47eb4c Dawsh 2y ago return query.all()
319 a47eb4c Dawsh 2y ago
320 a47eb4c Dawsh 2y ago
321 9ca312d Dawsh 2y ago def get_user_total_usage(
322 9ca312d Dawsh 2y ago db: Session, user: User, start: datetime, end: datetime, per_day=False
323 9ca312d Dawsh 2y ago ):
324 9ca312d Dawsh 2y ago usages = defaultdict(int)
325 9ca312d Dawsh 2y ago
326 9ca312d Dawsh 2y ago query = db.query(
327 9ca312d Dawsh 2y ago (
328 9ca312d Dawsh 2y ago cast(NodeUserUsage.created_at, Date).label("day")
329 9ca312d Dawsh 2y ago if per_day
330 9ca312d Dawsh 2y ago else NodeUserUsage.created_at
331 9ca312d Dawsh 2y ago ),
332 9ca312d Dawsh 2y ago func.sum(NodeUserUsage.used_traffic),
333 9ca312d Dawsh 2y ago ).filter(
334 9ca312d Dawsh 2y ago and_(
335 9ca312d Dawsh 2y ago NodeUserUsage.user_id == user.id,
336 9ca312d Dawsh 2y ago NodeUserUsage.created_at >= start,
337 9ca312d Dawsh 2y ago NodeUserUsage.created_at <= end,
338 9ca312d Dawsh 2y ago )
339 9ca312d Dawsh 2y ago )
340 9ca312d Dawsh 2y ago if per_day:
341 9ca312d Dawsh 2y ago query = query.group_by(cast(NodeUserUsage.created_at, Date))
342 9ca312d Dawsh 2y ago else:
343 9ca312d Dawsh 2y ago query = query.group_by(NodeUserUsage.created_at)
344 9ca312d Dawsh 2y ago
345 9ca312d Dawsh 2y ago for date, used_traffic in query:
346 9ca312d Dawsh 2y ago if per_day:
347 9ca312d Dawsh 2y ago timestamp = datetime(
348 9ca312d Dawsh 2y ago date.year, date.month, date.day, tzinfo=timezone.utc
349 9ca312d Dawsh 2y ago ).timestamp()
350 9ca312d Dawsh 2y ago usages[timestamp] += int(used_traffic)
351 9ca312d Dawsh 2y ago else:
352 9ca312d Dawsh 2y ago usages[date.replace(tzinfo=timezone.utc).timestamp()] += int(
353 9ca312d Dawsh 2y ago used_traffic
354 9ca312d Dawsh 2y ago )
355 9ca312d Dawsh 2y ago
356 9ca312d Dawsh 2y ago result = TrafficUsageSeries(usages=[])
357 9ca312d Dawsh 2y ago current = start.astimezone(timezone.utc).replace(
358 9ca312d Dawsh 2y ago minute=0, second=0, microsecond=0
359 9ca312d Dawsh 2y ago )
360 9ca312d Dawsh 2y ago if per_day:
361 9ca312d Dawsh 2y ago current = current.replace(hour=0)
362 9ca312d Dawsh 2y ago
363 9ca312d Dawsh 2y ago step = timedelta(days=1) if per_day else timedelta(hours=1)
364 9ca312d Dawsh 2y ago while current <= end.replace(tzinfo=timezone.utc):
365 9ca312d Dawsh 2y ago current_usage = usages.get(current.timestamp()) or 0
366 9ca312d Dawsh 2y ago result.usages.append((int(current.timestamp()), current_usage))
367 9ca312d Dawsh 2y ago result.total += current_usage
368 9ca312d Dawsh 2y ago current += step
369 9ca312d Dawsh 2y ago
370 9ca312d Dawsh 2y ago result.step = int(step.total_seconds())
371 9ca312d Dawsh 2y ago return result
372 9ca312d Dawsh 2y ago
373 9ca312d Dawsh 2y ago
374 6a4fee1 Dawsh 2y ago def get_total_usages(
375 6a4fee1 Dawsh 2y ago db: Session, admin: Admin, start: datetime, end: datetime
376 b15de9e Erfan 2y ago ) -> TrafficUsageSeries:
377 6a4fee1 Dawsh 2y ago usages = defaultdict(int)
378 6a4fee1 Dawsh 2y ago
379 6a4fee1 Dawsh 2y ago query = (
380 6a4fee1 Dawsh 2y ago db.query(
381 6a4fee1 Dawsh 2y ago NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
382 6a4fee1 Dawsh 2y ago )
383 6a4fee1 Dawsh 2y ago .group_by(NodeUserUsage.created_at)
384 6a4fee1 Dawsh 2y ago .filter(
385 6a4fee1 Dawsh 2y ago and_(
386 6a4fee1 Dawsh 2y ago NodeUserUsage.created_at >= start,
387 6a4fee1 Dawsh 2y ago NodeUserUsage.created_at <= end,
388 6a4fee1 Dawsh 2y ago )
389 6a4fee1 Dawsh 2y ago )
390 6a4fee1 Dawsh 2y ago )
391 6a4fee1 Dawsh 2y ago
392 6a4fee1 Dawsh 2y ago if not admin.is_sudo:
393 6a4fee1 Dawsh 2y ago query = (
394 6a4fee1 Dawsh 2y ago query.filter(
395 6a4fee1 Dawsh 2y ago Admin.id == admin.id,
396 6a4fee1 Dawsh 2y ago )
397 6a4fee1 Dawsh 2y ago .join(User, NodeUserUsage.user_id == User.id)
398 6a4fee1 Dawsh 2y ago .join(Admin, User.admin_id == Admin.id)
399 6a4fee1 Dawsh 2y ago )
400 b15de9e Erfan 2y ago
401 6a4fee1 Dawsh 2y ago for created_at, used_traffic in query.all():
402 b15de9e Erfan 2y ago timestamp = created_at.replace(tzinfo=timezone.utc).timestamp()
403 b15de9e Erfan 2y ago usages[timestamp] += int(used_traffic)
404 6a4fee1 Dawsh 2y ago
405 b15de9e Erfan 2y ago result = TrafficUsageSeries(usages=[], total=0)
406 6a4fee1 Dawsh 2y ago current = start.astimezone(timezone.utc).replace(
407 6a4fee1 Dawsh 2y ago minute=0, second=0, microsecond=0
408 6a4fee1 Dawsh 2y ago )
409 6a4fee1 Dawsh 2y ago
410 6a4fee1 Dawsh 2y ago while current <= end.replace(tzinfo=timezone.utc):
411 b15de9e Erfan 2y ago usage = usages.get(current.timestamp()) or 0
412 b15de9e Erfan 2y ago result.usages.append((int(current.timestamp()), usage))
413 b15de9e Erfan 2y ago result.total += usage
414 6a4fee1 Dawsh 2y ago current += timedelta(hours=1)
415 6a4fee1 Dawsh 2y ago
416 6a4fee1 Dawsh 2y ago return result
417 6a4fee1 Dawsh 2y ago
418 6a4fee1 Dawsh 2y ago
419 01cdf2e Dawsh 2y ago def get_user_usages(
420 01cdf2e Dawsh 2y ago db: Session,
421 e6e6467 Dawsh 2y ago db_user: User,
422 01cdf2e Dawsh 2y ago start: datetime,
423 01cdf2e Dawsh 2y ago end: datetime,
424 e6e6467 Dawsh 2y ago ) -> UserUsageSeriesResponse:
425 b15de9e Erfan 2y ago usages = defaultdict(dict)
426 b15de9e Erfan 2y ago
427 01cdf2e Dawsh 2y ago cond = and_(
428 e6e6467 Dawsh 2y ago NodeUserUsage.user_id == db_user.id,
429 01cdf2e Dawsh 2y ago NodeUserUsage.created_at >= start,
430 01cdf2e Dawsh 2y ago NodeUserUsage.created_at <= end,
431 01cdf2e Dawsh 2y ago )
432 a47eb4c Dawsh 2y ago
433 a47eb4c Dawsh 2y ago for v in db.query(NodeUserUsage).filter(cond):
434 b15de9e Erfan 2y ago timestamp = v.created_at.replace(tzinfo=timezone.utc).timestamp()
435 b15de9e Erfan 2y ago usages[v.node_id][timestamp] = v.used_traffic
436 e6e6467 Dawsh 2y ago
437 e6e6467 Dawsh 2y ago node_ids = list(usages.keys())
438 e6e6467 Dawsh 2y ago nodes = db.query(Node).where(Node.id.in_(node_ids))
439 e6e6467 Dawsh 2y ago node_id_names = {node.id: node.name for node in nodes}
440 e6e6467 Dawsh 2y ago
441 b15de9e Erfan 2y ago result = UserUsageSeriesResponse(
442 b15de9e Erfan 2y ago username=db_user.username, node_usages=[], total=0
443 b15de9e Erfan 2y ago )
444 b15de9e Erfan 2y ago
445 e6e6467 Dawsh 2y ago for node_id, rows in usages.items():
446 e6e6467 Dawsh 2y ago node_usages = UserNodeUsageSeries(
447 e6e6467 Dawsh 2y ago node_id=node_id, node_name=node_id_names[node_id], usages=[]
448 e6e6467 Dawsh 2y ago )
449 e6e6467 Dawsh 2y ago current = start.astimezone(timezone.utc).replace(
450 e6e6467 Dawsh 2y ago minute=0, second=0, microsecond=0
451 e6e6467 Dawsh 2y ago )
452 a47eb4c Dawsh 2y ago
453 e6e6467 Dawsh 2y ago while current <= end:
454 b15de9e Erfan 2y ago usage = rows.get(current.timestamp()) or 0
455 b15de9e Erfan 2y ago node_usages.usages.append((int(current.timestamp()), usage))
456 e6e6467 Dawsh 2y ago current += timedelta(hours=1)
457 b15de9e Erfan 2y ago result.total += usage
458 e6e6467 Dawsh 2y ago result.node_usages.append(node_usages)
459 b15de9e Erfan 2y ago
460 e6e6467 Dawsh 2y ago return result
461 a47eb4c Dawsh 2y ago
462 a47eb4c Dawsh 2y ago
463 e5dcd82 Dawsh 2y ago def get_users_count(
464 e5dcd82 Dawsh 2y ago db: Session,
465 e5dcd82 Dawsh 2y ago admin: Admin | None = None,
466 e5dcd82 Dawsh 2y ago enabled: bool | None = None,
467 e5dcd82 Dawsh 2y ago online: bool | None = None,
468 d377c4c Dawsh 2y ago expire_strategy: UserExpireStrategy | None = None,
469 d377c4c Dawsh 2y ago is_active: bool | None = None,
470 d377c4c Dawsh 2y ago expired: bool | None = None,
471 d377c4c Dawsh 2y ago data_limit_reached: bool | None = None,
472 e5dcd82 Dawsh 2y ago ):
473 45803f3 Dawsh 2y ago query = db.query(User.id).filter(User.removed == False)
474 a47eb4c Dawsh 2y ago if admin:
475 9ddfbc8 Mardin 2y ago query = query.filter(User.admin_id == admin.id)
476 d377c4c Dawsh 2y ago if is_active:
477 d377c4c Dawsh 2y ago query = query.filter(User.is_active == is_active)
478 d377c4c Dawsh 2y ago if expired:
479 d377c4c Dawsh 2y ago query = query.filter(User.expired == expired)
480 d377c4c Dawsh 2y ago if data_limit_reached:
481 d377c4c Dawsh 2y ago query = query.filter(User.data_limit_reached == data_limit_reached)
482 e5dcd82 Dawsh 2y ago if enabled:
483 e5dcd82 Dawsh 2y ago query = query.filter(User.enabled == enabled)
484 d377c4c Dawsh 2y ago if online is True:
485 e5dcd82 Dawsh 2y ago query = query.filter(
486 e5dcd82 Dawsh 2y ago User.online_at > (datetime.utcnow() - timedelta(seconds=30))
487 e5dcd82 Dawsh 2y ago )
488 d377c4c Dawsh 2y ago elif online is False:
489 d377c4c Dawsh 2y ago query = query.filter(
490 d377c4c Dawsh 2y ago User.online_at < (datetime.utcnow() - timedelta(seconds=30))
491 d377c4c Dawsh 2y ago )
492 d377c4c Dawsh 2y ago if expire_strategy:
493 d377c4c Dawsh 2y ago query = query.filter(User.expire_strategy == expire_strategy)
494 d377c4c Dawsh 2y ago
495 a47eb4c Dawsh 2y ago return query.count()
496 a47eb4c Dawsh 2y ago
497 a47eb4c Dawsh 2y ago
498 4708fcb Dawsh 2y ago def create_user(
499 4708fcb Dawsh 2y ago db: Session,
500 4708fcb Dawsh 2y ago user: UserCreate,
501 4708fcb Dawsh 2y ago admin: Admin = None,
502 4708fcb Dawsh 2y ago allowed_services: list | None = None,
503 4708fcb Dawsh 2y ago ):
504 4708fcb Dawsh 2y ago service_ids = (
505 4708fcb Dawsh 2y ago [sid for sid in user.service_ids if sid in allowed_services]
506 4708fcb Dawsh 2y ago if allowed_services is not None
507 4708fcb Dawsh 2y ago else user.service_ids
508 4708fcb Dawsh 2y ago )
509 a47eb4c Dawsh 2y ago dbuser = User(
510 a47eb4c Dawsh 2y ago username=user.username,
511 a47eb4c Dawsh 2y ago key=user.key,
512 d377c4c Dawsh 2y ago expire_strategy=user.expire_strategy,
513 d377c4c Dawsh 2y ago expire_date=user.expire_date,
514 d377c4c Dawsh 2y ago usage_duration=user.usage_duration,
515 d377c4c Dawsh 2y ago activation_deadline=user.activation_deadline,
516 01cdf2e Dawsh 2y ago services=db.query(Service)
517 4708fcb Dawsh 2y ago .filter(Service.id.in_(service_ids))
518 01cdf2e Dawsh 2y ago .all(), # user.services,
519 a47eb4c Dawsh 2y ago data_limit=(user.data_limit or None),
520 a47eb4c Dawsh 2y ago admin=admin,
521 a47eb4c Dawsh 2y ago data_limit_reset_strategy=user.data_limit_reset_strategy,
522 a47eb4c Dawsh 2y ago note=user.note,
523 a47eb4c Dawsh 2y ago )
524 a47eb4c Dawsh 2y ago db.add(dbuser)
525 a47eb4c Dawsh 2y ago db.commit()
526 a47eb4c Dawsh 2y ago db.refresh(dbuser)
527 a47eb4c Dawsh 2y ago return dbuser
528 a47eb4c Dawsh 2y ago
529 a47eb4c Dawsh 2y ago
530 a47eb4c Dawsh 2y ago def remove_user(db: Session, dbuser: User):
531 fdd3f61 Dawsh 2y ago dbuser.username = None
532 24826a2 Dawsh 2y ago dbuser.removed = True
533 24826a2 Dawsh 2y ago dbuser.activated = False
534 a47eb4c Dawsh 2y ago # db.query(User).filter_by(id=user_id).delete()
535 a47eb4c Dawsh 2y ago db.commit()
536 a47eb4c Dawsh 2y ago
537 a47eb4c Dawsh 2y ago
538 4708fcb Dawsh 2y ago def update_user(
539 4708fcb Dawsh 2y ago db: Session,
540 4708fcb Dawsh 2y ago dbuser: User,
541 4708fcb Dawsh 2y ago modify: UserModify,
542 4708fcb Dawsh 2y ago allowed_services: list | None = None,
543 4708fcb Dawsh 2y ago ):
544 a47eb4c Dawsh 2y ago if modify.data_limit is not None:
545 d377c4c Dawsh 2y ago dbuser.data_limit = modify.data_limit or None
546 d377c4c Dawsh 2y ago
547 d377c4c Dawsh 2y ago if modify.expire_strategy is not None:
548 d377c4c Dawsh 2y ago dbuser.expire_strategy = modify.expire_strategy or None
549 d377c4c Dawsh 2y ago
550 4442e87 mahdim10 2y ago if modify.expire_strategy == UserExpireStrategy.FIXED_DATE:
551 4442e87 mahdim10 2y ago dbuser.usage_duration = None
552 4442e87 mahdim10 2y ago dbuser.activation_deadline = None
553 4442e87 mahdim10 2y ago elif modify.expire_strategy == UserExpireStrategy.START_ON_FIRST_USE:
554 4442e87 mahdim10 2y ago dbuser.expire_date = None
555 4442e87 mahdim10 2y ago elif modify.expire_strategy == UserExpireStrategy.NEVER:
556 4442e87 mahdim10 2y ago dbuser.expire_date = None
557 4442e87 mahdim10 2y ago dbuser.usage_duration = None
558 4442e87 mahdim10 2y ago dbuser.activation_deadline = None
559 4442e87 mahdim10 2y ago
560 d377c4c Dawsh 2y ago if modify.expire_date is not None:
561 d377c4c Dawsh 2y ago dbuser.expire_date = modify.expire_date or None
562 a47eb4c Dawsh 2y ago
563 a47eb4c Dawsh 2y ago if modify.note is not None:
564 a47eb4c Dawsh 2y ago dbuser.note = modify.note or None
565 a47eb4c Dawsh 2y ago
566 a47eb4c Dawsh 2y ago if modify.data_limit_reset_strategy is not None:
567 d377c4c Dawsh 2y ago dbuser.data_limit_reset_strategy = modify.data_limit_reset_strategy
568 a47eb4c Dawsh 2y ago
569 d377c4c Dawsh 2y ago if modify.activation_deadline is not None:
570 d377c4c Dawsh 2y ago dbuser.activation_deadline = modify.activation_deadline
571 a47eb4c Dawsh 2y ago
572 d377c4c Dawsh 2y ago if modify.usage_duration is not None:
573 d377c4c Dawsh 2y ago dbuser.usage_duration = modify.usage_duration
574 2de9d72 Dawsh 2y ago
575 db68a58 Dawsh 2y ago if modify.service_ids is not None:
576 4708fcb Dawsh 2y ago if allowed_services is not None:
577 4708fcb Dawsh 2y ago service_ids = [
578 4708fcb Dawsh 2y ago sid for sid in modify.service_ids if sid in allowed_services
579 4708fcb Dawsh 2y ago ]
580 4708fcb Dawsh 2y ago else:
581 4708fcb Dawsh 2y ago service_ids = modify.service_ids
582 4708fcb Dawsh 2y ago
583 01cdf2e Dawsh 2y ago dbuser.services = (
584 4708fcb Dawsh 2y ago db.query(Service).filter(Service.id.in_(service_ids)).all()
585 01cdf2e Dawsh 2y ago )
586 a47eb4c Dawsh 2y ago dbuser.edit_at = datetime.utcnow()
587 a47eb4c Dawsh 2y ago
588 a47eb4c Dawsh 2y ago db.commit()
589 a47eb4c Dawsh 2y ago db.refresh(dbuser)
590 a47eb4c Dawsh 2y ago return dbuser
591 a47eb4c Dawsh 2y ago
592 a47eb4c Dawsh 2y ago
593 a47eb4c Dawsh 2y ago def reset_user_data_usage(db: Session, dbuser: User):
594 42a1073 Dawsh 2y ago dbuser.traffic_reset_at = datetime.utcnow()
595 a47eb4c Dawsh 2y ago
596 a47eb4c Dawsh 2y ago dbuser.used_traffic = 0
597 42a1073 Dawsh 2y ago
598 a47eb4c Dawsh 2y ago db.add(dbuser)
599 a47eb4c Dawsh 2y ago
600 a47eb4c Dawsh 2y ago db.commit()
601 a47eb4c Dawsh 2y ago db.refresh(dbuser)
602 a47eb4c Dawsh 2y ago return dbuser
603 a47eb4c Dawsh 2y ago
604 a47eb4c Dawsh 2y ago
605 a47eb4c Dawsh 2y ago def revoke_user_sub(db: Session, dbuser: User):
606 34a7612 Dawsh 2y ago dbuser.key = secrets.token_hex(16)
607 00d6be5 Erfan 2y ago dbuser.sub_revoked_at = datetime.utcnow()
608 a47eb4c Dawsh 2y ago db.commit()
609 a47eb4c Dawsh 2y ago db.refresh(dbuser)
610 a47eb4c Dawsh 2y ago return dbuser
611 a47eb4c Dawsh 2y ago
612 a47eb4c Dawsh 2y ago
613 a47eb4c Dawsh 2y ago def update_user_sub(db: Session, dbuser: User, user_agent: str):
614 a47eb4c Dawsh 2y ago dbuser.sub_updated_at = datetime.utcnow()
615 a47eb4c Dawsh 2y ago dbuser.sub_last_user_agent = user_agent
616 a47eb4c Dawsh 2y ago
617 a47eb4c Dawsh 2y ago db.commit()
618 a47eb4c Dawsh 2y ago db.refresh(dbuser)
619 a47eb4c Dawsh 2y ago return dbuser
620 a47eb4c Dawsh 2y ago
621 a47eb4c Dawsh 2y ago
622 a47eb4c Dawsh 2y ago def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
623 a47eb4c Dawsh 2y ago query = db.query(User)
624 a47eb4c Dawsh 2y ago
625 a47eb4c Dawsh 2y ago if admin:
626 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
627 a47eb4c Dawsh 2y ago
628 d377c4c Dawsh 2y ago for db_user in query.all():
629 d377c4c Dawsh 2y ago db_user.used_traffic = 0
630 a47eb4c Dawsh 2y ago
631 a47eb4c Dawsh 2y ago db.commit()
632 a47eb4c Dawsh 2y ago
633 a47eb4c Dawsh 2y ago
634 a47eb4c Dawsh 2y ago def update_user_status(db: Session, dbuser: User, status: UserStatus):
635 a47eb4c Dawsh 2y ago dbuser.status = status
636 a47eb4c Dawsh 2y ago db.commit()
637 a47eb4c Dawsh 2y ago db.refresh(dbuser)
638 a47eb4c Dawsh 2y ago return dbuser
639 a47eb4c Dawsh 2y ago
640 a47eb4c Dawsh 2y ago
641 a47eb4c Dawsh 2y ago def set_owner(db: Session, dbuser: User, admin: Admin):
642 a47eb4c Dawsh 2y ago dbuser.admin = admin
643 a47eb4c Dawsh 2y ago db.commit()
644 a47eb4c Dawsh 2y ago db.refresh(dbuser)
645 a47eb4c Dawsh 2y ago return dbuser
646 a47eb4c Dawsh 2y ago
647 a47eb4c Dawsh 2y ago
648 a47eb4c Dawsh 2y ago def get_system_usage(db: Session):
649 a47eb4c Dawsh 2y ago return db.query(System).first()
650 a47eb4c Dawsh 2y ago
651 a47eb4c Dawsh 2y ago
652 a47eb4c Dawsh 2y ago def get_jwt_secret_key(db: Session):
653 a47eb4c Dawsh 2y ago return db.query(JWT).first().secret_key
654 a47eb4c Dawsh 2y ago
655 a47eb4c Dawsh 2y ago
656 a47eb4c Dawsh 2y ago def get_tls_certificate(db: Session):
657 a47eb4c Dawsh 2y ago return db.query(TLS).first()
658 a47eb4c Dawsh 2y ago
659 a47eb4c Dawsh 2y ago
660 4708fcb Dawsh 2y ago def get_admin(db: Session, username: str) -> Admin | None:
661 a47eb4c Dawsh 2y ago return db.query(Admin).filter(Admin.username == username).first()
662 a47eb4c Dawsh 2y ago
663 a47eb4c Dawsh 2y ago
664 a47eb4c Dawsh 2y ago def create_admin(db: Session, admin: AdminCreate):
665 a47eb4c Dawsh 2y ago dbadmin = Admin(
666 a47eb4c Dawsh 2y ago username=admin.username,
667 a47eb4c Dawsh 2y ago hashed_password=admin.hashed_password,
668 01cdf2e Dawsh 2y ago is_sudo=admin.is_sudo,
669 4708fcb Dawsh 2y ago enabled=admin.enabled,
670 4708fcb Dawsh 2y ago all_services_access=admin.all_services_access,
671 4708fcb Dawsh 2y ago modify_users_access=admin.modify_users_access,
672 4708fcb Dawsh 2y ago services=db.query(Service)
673 4708fcb Dawsh 2y ago .filter(Service.id.in_(admin.service_ids))
674 4708fcb Dawsh 2y ago .all(),
675 4708fcb Dawsh 2y ago subscription_url_prefix=admin.subscription_url_prefix,
676 a47eb4c Dawsh 2y ago )
677 a47eb4c Dawsh 2y ago db.add(dbadmin)
678 a47eb4c Dawsh 2y ago db.commit()
679 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
680 a47eb4c Dawsh 2y ago return dbadmin
681 a47eb4c Dawsh 2y ago
682 a47eb4c Dawsh 2y ago
683 4708fcb Dawsh 2y ago def update_admin(
684 4708fcb Dawsh 2y ago db: Session, dbadmin: Admin, modifications: AdminPartialModify
685 4708fcb Dawsh 2y ago ):
686 4708fcb Dawsh 2y ago for attribute in [
687 4708fcb Dawsh 2y ago "is_sudo",
688 4708fcb Dawsh 2y ago "hashed_password",
689 4708fcb Dawsh 2y ago "enabled",
690 4708fcb Dawsh 2y ago "all_services_access",
691 4708fcb Dawsh 2y ago "modify_users_access",
692 4708fcb Dawsh 2y ago "subscription_url_prefix",
693 4708fcb Dawsh 2y ago ]:
694 4708fcb Dawsh 2y ago if not isinstance(getattr(modifications, attribute), NoneType):
695 4708fcb Dawsh 2y ago setattr(dbadmin, attribute, getattr(modifications, attribute))
696 4708fcb Dawsh 2y ago if attribute == "hashed_password":
697 4708fcb Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
698 4708fcb Dawsh 2y ago if isinstance(modifications.service_ids, list):
699 4708fcb Dawsh 2y ago dbadmin.services = (
700 4708fcb Dawsh 2y ago db.query(Service)
701 4708fcb Dawsh 2y ago .filter(Service.id.in_(modifications.service_ids))
702 4708fcb Dawsh 2y ago .all()
703 4708fcb Dawsh 2y ago )
704 a47eb4c Dawsh 2y ago db.commit()
705 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
706 2de9d72 Dawsh 2y ago return dbadmin
707 a47eb4c Dawsh 2y ago
708 a47eb4c Dawsh 2y ago
709 01cdf2e Dawsh 2y ago def partial_update_admin(
710 01cdf2e Dawsh 2y ago db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
711 01cdf2e Dawsh 2y ago ):
712 a47eb4c Dawsh 2y ago if modified_admin.is_sudo is not None:
713 a47eb4c Dawsh 2y ago dbadmin.is_sudo = modified_admin.is_sudo
714 01cdf2e Dawsh 2y ago if (
715 01cdf2e Dawsh 2y ago modified_admin.password is not None
716 01cdf2e Dawsh 2y ago and dbadmin.hashed_password != modified_admin.hashed_password
717 01cdf2e Dawsh 2y ago ):
718 a47eb4c Dawsh 2y ago dbadmin.hashed_password = modified_admin.hashed_password
719 a47eb4c Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
720 a47eb4c Dawsh 2y ago
721 a47eb4c Dawsh 2y ago db.commit()
722 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
723 a47eb4c Dawsh 2y ago return dbadmin
724 a47eb4c Dawsh 2y ago
725 a47eb4c Dawsh 2y ago
726 a47eb4c Dawsh 2y ago def remove_admin(db: Session, dbadmin: Admin):
727 a47eb4c Dawsh 2y ago db.delete(dbadmin)
728 a47eb4c Dawsh 2y ago db.commit()
729 a47eb4c Dawsh 2y ago return dbadmin
730 a47eb4c Dawsh 2y ago
731 a47eb4c Dawsh 2y ago
732 01cdf2e Dawsh 2y ago def get_admins(
733 01cdf2e Dawsh 2y ago db: Session,
734 01cdf2e Dawsh 2y ago offset: Optional[int] = None,
735 01cdf2e Dawsh 2y ago limit: Optional[int] = None,
736 01cdf2e Dawsh 2y ago username: Optional[str] = None,
737 01cdf2e Dawsh 2y ago ):
738 a47eb4c Dawsh 2y ago query = db.query(Admin)
739 a47eb4c Dawsh 2y ago if username:
740 01cdf2e Dawsh 2y ago query = query.filter(Admin.username.ilike(f"%{username}%"))
741 a47eb4c Dawsh 2y ago if offset:
742 a47eb4c Dawsh 2y ago query = query.offset(offset)
743 a47eb4c Dawsh 2y ago if limit:
744 a47eb4c Dawsh 2y ago query = query.limit(limit)
745 a47eb4c Dawsh 2y ago return query.all()
746 a47eb4c Dawsh 2y ago
747 a47eb4c Dawsh 2y ago
748 ee1351a Dawsh 2y ago def create_service(db: Session, service: ServiceCreate) -> Service:
749 a47eb4c Dawsh 2y ago dbservice = Service(
750 a47eb4c Dawsh 2y ago name=service.name,
751 f1ab16c Dawsh 2y ago inbounds=db.query(Inbound)
752 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(service.inbound_ids))
753 f1ab16c Dawsh 2y ago .all(),
754 db68a58 Dawsh 2y ago users=[],
755 2de9d72 Dawsh 2y ago )
756 a47eb4c Dawsh 2y ago db.add(dbservice)
757 a47eb4c Dawsh 2y ago db.commit()
758 a47eb4c Dawsh 2y ago db.refresh(dbservice)
759 a47eb4c Dawsh 2y ago return dbservice
760 a47eb4c Dawsh 2y ago
761 a47eb4c Dawsh 2y ago
762 ee1351a Dawsh 2y ago def get_service(db: Session, service_id: id) -> Service:
763 ee1351a Dawsh 2y ago return db.query(Service).filter(Service.id == service_id).first()
764 a47eb4c Dawsh 2y ago
765 12b255b Dawsh 2y ago
766 a47eb4c Dawsh 2y ago def get_services(db: Session) -> List[Service]:
767 a47eb4c Dawsh 2y ago return db.query(Service).all()
768 a47eb4c Dawsh 2y ago
769 12b255b Dawsh 2y ago
770 f1ab16c Dawsh 2y ago def update_service(
771 f1ab16c Dawsh 2y ago db: Session, db_service: Service, modification: ServiceModify
772 f1ab16c Dawsh 2y ago ):
773 ee1351a Dawsh 2y ago if modification.name is not None:
774 ee1351a Dawsh 2y ago db_service.name = modification.name
775 ee1351a Dawsh 2y ago
776 db68a58 Dawsh 2y ago if modification.inbound_ids is not None:
777 01cdf2e Dawsh 2y ago db_service.inbounds = (
778 f1ab16c Dawsh 2y ago db.query(Inbound)
779 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(modification.inbound_ids))
780 f1ab16c Dawsh 2y ago .all()
781 01cdf2e Dawsh 2y ago )
782 ee1351a Dawsh 2y ago
783 ee1351a Dawsh 2y ago db.commit()
784 ee1351a Dawsh 2y ago db.refresh(db_service)
785 ee1351a Dawsh 2y ago return db_service
786 ee1351a Dawsh 2y ago
787 ee1351a Dawsh 2y ago
788 ee1351a Dawsh 2y ago def remove_service(db: Session, db_service: Service):
789 ee1351a Dawsh 2y ago db.delete(db_service)
790 ee1351a Dawsh 2y ago db.commit()
791 ee1351a Dawsh 2y ago return db_service
792 ee1351a Dawsh 2y ago
793 ee1351a Dawsh 2y ago
794 a47eb4c Dawsh 2y ago def get_node(db: Session, name: str):
795 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.name == name).first()
796 a47eb4c Dawsh 2y ago
797 a47eb4c Dawsh 2y ago
798 a47eb4c Dawsh 2y ago def get_node_by_id(db: Session, node_id: int):
799 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.id == node_id).first()
800 a47eb4c Dawsh 2y ago
801 a47eb4c Dawsh 2y ago
802 01cdf2e Dawsh 2y ago def get_nodes(
803 f1ab16c Dawsh 2y ago db: Session,
804 f1ab16c Dawsh 2y ago status: Optional[Union[NodeStatus, list]] = None,
805 f1ab16c Dawsh 2y ago enabled: bool = None,
806 01cdf2e Dawsh 2y ago ):
807 a47eb4c Dawsh 2y ago query = db.query(Node)
808 a47eb4c Dawsh 2y ago
809 a47eb4c Dawsh 2y ago if status:
810 a47eb4c Dawsh 2y ago if isinstance(status, list):
811 a47eb4c Dawsh 2y ago query = query.filter(Node.status.in_(status))
812 a47eb4c Dawsh 2y ago else:
813 a47eb4c Dawsh 2y ago query = query.filter(Node.status == status)
814 a47eb4c Dawsh 2y ago
815 a47eb4c Dawsh 2y ago if enabled:
816 a47eb4c Dawsh 2y ago query = query.filter(Node.status != NodeStatus.disabled)
817 a47eb4c Dawsh 2y ago
818 a47eb4c Dawsh 2y ago return query.all()
819 a47eb4c Dawsh 2y ago
820 a47eb4c Dawsh 2y ago
821 78a3e64 Dawsh 2y ago def get_node_usage(
822 78a3e64 Dawsh 2y ago db: Session, start: datetime, end: datetime, node: Node
823 78a3e64 Dawsh 2y ago ) -> TrafficUsageSeries:
824 78a3e64 Dawsh 2y ago usages = defaultdict(int)
825 a47eb4c Dawsh 2y ago
826 78a3e64 Dawsh 2y ago query = (
827 78a3e64 Dawsh 2y ago db.query(
828 78a3e64 Dawsh 2y ago NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
829 78a3e64 Dawsh 2y ago )
830 78a3e64 Dawsh 2y ago .group_by(NodeUserUsage.created_at)
831 78a3e64 Dawsh 2y ago .filter(
832 78a3e64 Dawsh 2y ago and_(
833 78a3e64 Dawsh 2y ago NodeUserUsage.node_id == node.id,
834 78a3e64 Dawsh 2y ago NodeUserUsage.created_at >= start,
835 78a3e64 Dawsh 2y ago NodeUserUsage.created_at <= end,
836 78a3e64 Dawsh 2y ago )
837 78a3e64 Dawsh 2y ago )
838 a47eb4c Dawsh 2y ago )
839 78a3e64 Dawsh 2y ago
840 78a3e64 Dawsh 2y ago for created_at, used_traffic in query.all():
841 78a3e64 Dawsh 2y ago usages[created_at.replace(tzinfo=timezone.utc).timestamp()] += int(
842 78a3e64 Dawsh 2y ago used_traffic
843 a47eb4c Dawsh 2y ago )
844 a47eb4c Dawsh 2y ago
845 78a3e64 Dawsh 2y ago result = TrafficUsageSeries(usages=[], total=0)
846 78a3e64 Dawsh 2y ago current = start.astimezone(timezone.utc).replace(
847 78a3e64 Dawsh 2y ago minute=0, second=0, microsecond=0
848 78a3e64 Dawsh 2y ago )
849 a47eb4c Dawsh 2y ago
850 78a3e64 Dawsh 2y ago while current <= end.replace(tzinfo=timezone.utc):
851 78a3e64 Dawsh 2y ago usage = usages.get(current.timestamp()) or 0
852 b4ec790 Dawsh 2y ago result.usages.append((int(current.timestamp()), usage))
853 78a3e64 Dawsh 2y ago result.total += usage
854 78a3e64 Dawsh 2y ago current += timedelta(hours=1)
855 a47eb4c Dawsh 2y ago
856 78a3e64 Dawsh 2y ago return result
857 a47eb4c Dawsh 2y ago
858 a47eb4c Dawsh 2y ago
859 a47eb4c Dawsh 2y ago def create_node(db: Session, node: NodeCreate):
860 01cdf2e Dawsh 2y ago dbnode = Node(
861 01cdf2e Dawsh 2y ago name=node.name,
862 01cdf2e Dawsh 2y ago address=node.address,
863 01cdf2e Dawsh 2y ago port=node.port,
864 01cdf2e Dawsh 2y ago connection_backend=node.connection_backend,
865 01cdf2e Dawsh 2y ago )
866 a47eb4c Dawsh 2y ago
867 a47eb4c Dawsh 2y ago db.add(dbnode)
868 a47eb4c Dawsh 2y ago db.commit()
869 a47eb4c Dawsh 2y ago db.refresh(dbnode)
870 a47eb4c Dawsh 2y ago return dbnode
871 a47eb4c Dawsh 2y ago
872 a47eb4c Dawsh 2y ago
873 a47eb4c Dawsh 2y ago def remove_node(db: Session, dbnode: Node):
874 a47eb4c Dawsh 2y ago db.delete(dbnode)
875 a47eb4c Dawsh 2y ago db.commit()
876 a47eb4c Dawsh 2y ago return dbnode
877 a47eb4c Dawsh 2y ago
878 a47eb4c Dawsh 2y ago
879 a47eb4c Dawsh 2y ago def update_node(db: Session, dbnode: Node, modify: NodeModify):
880 a47eb4c Dawsh 2y ago if modify.name is not None:
881 a47eb4c Dawsh 2y ago dbnode.name = modify.name
882 a47eb4c Dawsh 2y ago
883 a47eb4c Dawsh 2y ago if modify.address is not None:
884 a47eb4c Dawsh 2y ago dbnode.address = modify.address
885 a47eb4c Dawsh 2y ago
886 a47eb4c Dawsh 2y ago if modify.port is not None:
887 a47eb4c Dawsh 2y ago dbnode.port = modify.port
888 a47eb4c Dawsh 2y ago
889 a47eb4c Dawsh 2y ago if modify.status is NodeStatus.disabled:
890 a47eb4c Dawsh 2y ago dbnode.status = modify.status
891 a47eb4c Dawsh 2y ago dbnode.xray_version = None
892 a47eb4c Dawsh 2y ago dbnode.message = None
893 a47eb4c Dawsh 2y ago else:
894 12b255b Dawsh 2y ago dbnode.status = NodeStatus.unhealthy
895 a47eb4c Dawsh 2y ago
896 a47eb4c Dawsh 2y ago if modify.usage_coefficient:
897 a47eb4c Dawsh 2y ago dbnode.usage_coefficient = modify.usage_coefficient
898 a47eb4c Dawsh 2y ago
899 5584d4a Mrdn 2y ago if modify.connection_backend:
900 5584d4a Mrdn 2y ago dbnode.connection_backend = modify.connection_backend
901 5584d4a Mrdn 2y ago
902 a47eb4c Dawsh 2y ago db.commit()
903 a47eb4c Dawsh 2y ago db.refresh(dbnode)
904 a47eb4c Dawsh 2y ago return dbnode
905 a47eb4c Dawsh 2y ago
906 a47eb4c Dawsh 2y ago
907 01cdf2e Dawsh 2y ago def update_node_status(
908 01cdf2e Dawsh 2y ago db: Session,
909 01cdf2e Dawsh 2y ago node_id: int,
910 01cdf2e Dawsh 2y ago status: NodeStatus,
911 01cdf2e Dawsh 2y ago message: str = None,
912 01cdf2e Dawsh 2y ago version: str = None,
913 01cdf2e Dawsh 2y ago ):
914 f6a32ea Dawsh 2y ago db_node = db.query(Node).where(Node.id == node_id).first()
915 f6a32ea Dawsh 2y ago db_node.status = status
916 f6a32ea Dawsh 2y ago if message:
917 f6a32ea Dawsh 2y ago db_node.message = message
918 f6a32ea Dawsh 2y ago db_node.last_status_change = datetime.utcnow()
919 a47eb4c Dawsh 2y ago db.commit()