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