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 · nightly
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 ae5301f Erfan 1y ago with db.get_bind().begin() as conn:
729 ae5301f Erfan 1y ago conn.execute(
730 ae5301f Erfan 1y ago update(User)
731 ae5301f Erfan 1y ago .where(User.id == dbuser.id)
732 ae5301f Erfan 1y ago .values(
733 ae5301f Erfan 1y ago sub_updated_at=datetime.utcnow(),
734 ae5301f Erfan 1y ago sub_last_user_agent=user_agent,
735 ae5301f Erfan 1y ago )
736 ae5301f Erfan 1y ago )
737 a47eb4c Dawsh 2y ago
738 a47eb4c Dawsh 2y ago
739 a47eb4c Dawsh 2y ago def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
740 a47eb4c Dawsh 2y ago query = db.query(User)
741 a47eb4c Dawsh 2y ago
742 a47eb4c Dawsh 2y ago if admin:
743 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
744 a47eb4c Dawsh 2y ago
745 d377c4c Dawsh 2y ago for db_user in query.all():
746 d377c4c Dawsh 2y ago db_user.used_traffic = 0
747 a47eb4c Dawsh 2y ago
748 a47eb4c Dawsh 2y ago db.commit()
749 a47eb4c Dawsh 2y ago
750 a47eb4c Dawsh 2y ago
751 a47eb4c Dawsh 2y ago def update_user_status(db: Session, dbuser: User, status: UserStatus):
752 a47eb4c Dawsh 2y ago dbuser.status = status
753 a47eb4c Dawsh 2y ago db.commit()
754 a47eb4c Dawsh 2y ago db.refresh(dbuser)
755 a47eb4c Dawsh 2y ago return dbuser
756 a47eb4c Dawsh 2y ago
757 a47eb4c Dawsh 2y ago
758 a47eb4c Dawsh 2y ago def set_owner(db: Session, dbuser: User, admin: Admin):
759 a47eb4c Dawsh 2y ago dbuser.admin = admin
760 a47eb4c Dawsh 2y ago db.commit()
761 a47eb4c Dawsh 2y ago db.refresh(dbuser)
762 a47eb4c Dawsh 2y ago return dbuser
763 a47eb4c Dawsh 2y ago
764 a47eb4c Dawsh 2y ago
765 a47eb4c Dawsh 2y ago def get_system_usage(db: Session):
766 a47eb4c Dawsh 2y ago return db.query(System).first()
767 a47eb4c Dawsh 2y ago
768 a47eb4c Dawsh 2y ago
769 a47eb4c Dawsh 2y ago def get_jwt_secret_key(db: Session):
770 a47eb4c Dawsh 2y ago return db.query(JWT).first().secret_key
771 a47eb4c Dawsh 2y ago
772 a47eb4c Dawsh 2y ago
773 a47eb4c Dawsh 2y ago def get_tls_certificate(db: Session):
774 a47eb4c Dawsh 2y ago return db.query(TLS).first()
775 a47eb4c Dawsh 2y ago
776 a47eb4c Dawsh 2y ago
777 4708fcb Dawsh 2y ago def get_admin(db: Session, username: str) -> Admin | None:
778 a47eb4c Dawsh 2y ago return db.query(Admin).filter(Admin.username == username).first()
779 a47eb4c Dawsh 2y ago
780 a47eb4c Dawsh 2y ago
781 a47eb4c Dawsh 2y ago def create_admin(db: Session, admin: AdminCreate):
782 a47eb4c Dawsh 2y ago dbadmin = Admin(
783 a47eb4c Dawsh 2y ago username=admin.username,
784 a47eb4c Dawsh 2y ago hashed_password=admin.hashed_password,
785 01cdf2e Dawsh 2y ago is_sudo=admin.is_sudo,
786 4708fcb Dawsh 2y ago enabled=admin.enabled,
787 4708fcb Dawsh 2y ago all_services_access=admin.all_services_access,
788 4708fcb Dawsh 2y ago modify_users_access=admin.modify_users_access,
789 4708fcb Dawsh 2y ago services=db.query(Service)
790 4708fcb Dawsh 2y ago .filter(Service.id.in_(admin.service_ids))
791 4708fcb Dawsh 2y ago .all(),
792 4708fcb Dawsh 2y ago subscription_url_prefix=admin.subscription_url_prefix,
793 a47eb4c Dawsh 2y ago )
794 a47eb4c Dawsh 2y ago db.add(dbadmin)
795 a47eb4c Dawsh 2y ago db.commit()
796 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
797 a47eb4c Dawsh 2y ago return dbadmin
798 a47eb4c Dawsh 2y ago
799 a47eb4c Dawsh 2y ago
800 4708fcb Dawsh 2y ago def update_admin(
801 4708fcb Dawsh 2y ago db: Session, dbadmin: Admin, modifications: AdminPartialModify
802 4708fcb Dawsh 2y ago ):
803 4708fcb Dawsh 2y ago for attribute in [
804 4708fcb Dawsh 2y ago "is_sudo",
805 4708fcb Dawsh 2y ago "hashed_password",
806 4708fcb Dawsh 2y ago "enabled",
807 4708fcb Dawsh 2y ago "all_services_access",
808 4708fcb Dawsh 2y ago "modify_users_access",
809 4708fcb Dawsh 2y ago "subscription_url_prefix",
810 4708fcb Dawsh 2y ago ]:
811 4708fcb Dawsh 2y ago if not isinstance(getattr(modifications, attribute), NoneType):
812 4708fcb Dawsh 2y ago setattr(dbadmin, attribute, getattr(modifications, attribute))
813 4708fcb Dawsh 2y ago if attribute == "hashed_password":
814 4708fcb Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
815 4708fcb Dawsh 2y ago if isinstance(modifications.service_ids, list):
816 4708fcb Dawsh 2y ago dbadmin.services = (
817 4708fcb Dawsh 2y ago db.query(Service)
818 4708fcb Dawsh 2y ago .filter(Service.id.in_(modifications.service_ids))
819 4708fcb Dawsh 2y ago .all()
820 4708fcb Dawsh 2y ago )
821 a47eb4c Dawsh 2y ago db.commit()
822 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
823 2de9d72 Dawsh 2y ago return dbadmin
824 a47eb4c Dawsh 2y ago
825 a47eb4c Dawsh 2y ago
826 01cdf2e Dawsh 2y ago def partial_update_admin(
827 01cdf2e Dawsh 2y ago db: Session, dbadmin: Admin, modified_admin: AdminPartialModify
828 01cdf2e Dawsh 2y ago ):
829 a47eb4c Dawsh 2y ago if modified_admin.is_sudo is not None:
830 a47eb4c Dawsh 2y ago dbadmin.is_sudo = modified_admin.is_sudo
831 01cdf2e Dawsh 2y ago if (
832 01cdf2e Dawsh 2y ago modified_admin.password is not None
833 01cdf2e Dawsh 2y ago and dbadmin.hashed_password != modified_admin.hashed_password
834 01cdf2e Dawsh 2y ago ):
835 a47eb4c Dawsh 2y ago dbadmin.hashed_password = modified_admin.hashed_password
836 a47eb4c Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
837 a47eb4c Dawsh 2y ago
838 a47eb4c Dawsh 2y ago db.commit()
839 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
840 a47eb4c Dawsh 2y ago return dbadmin
841 a47eb4c Dawsh 2y ago
842 a47eb4c Dawsh 2y ago
843 a47eb4c Dawsh 2y ago def remove_admin(db: Session, dbadmin: Admin):
844 a47eb4c Dawsh 2y ago db.delete(dbadmin)
845 a47eb4c Dawsh 2y ago db.commit()
846 a47eb4c Dawsh 2y ago return dbadmin
847 a47eb4c Dawsh 2y ago
848 a47eb4c Dawsh 2y ago
849 01cdf2e Dawsh 2y ago def get_admins(
850 01cdf2e Dawsh 2y ago db: Session,
851 01cdf2e Dawsh 2y ago offset: Optional[int] = None,
852 01cdf2e Dawsh 2y ago limit: Optional[int] = None,
853 01cdf2e Dawsh 2y ago username: Optional[str] = None,
854 01cdf2e Dawsh 2y ago ):
855 a47eb4c Dawsh 2y ago query = db.query(Admin)
856 a47eb4c Dawsh 2y ago if username:
857 01cdf2e Dawsh 2y ago query = query.filter(Admin.username.ilike(f"%{username}%"))
858 a47eb4c Dawsh 2y ago if offset:
859 a47eb4c Dawsh 2y ago query = query.offset(offset)
860 a47eb4c Dawsh 2y ago if limit:
861 a47eb4c Dawsh 2y ago query = query.limit(limit)
862 a47eb4c Dawsh 2y ago return query.all()
863 a47eb4c Dawsh 2y ago
864 a47eb4c Dawsh 2y ago
865 ee1351a Dawsh 2y ago def create_service(db: Session, service: ServiceCreate) -> Service:
866 a47eb4c Dawsh 2y ago dbservice = Service(
867 a47eb4c Dawsh 2y ago name=service.name,
868 f1ab16c Dawsh 2y ago inbounds=db.query(Inbound)
869 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(service.inbound_ids))
870 f1ab16c Dawsh 2y ago .all(),
871 db68a58 Dawsh 2y ago users=[],
872 2de9d72 Dawsh 2y ago )
873 a47eb4c Dawsh 2y ago db.add(dbservice)
874 a47eb4c Dawsh 2y ago db.commit()
875 a47eb4c Dawsh 2y ago db.refresh(dbservice)
876 a47eb4c Dawsh 2y ago return dbservice
877 a47eb4c Dawsh 2y ago
878 a47eb4c Dawsh 2y ago
879 ee1351a Dawsh 2y ago def get_service(db: Session, service_id: id) -> Service:
880 ee1351a Dawsh 2y ago return db.query(Service).filter(Service.id == service_id).first()
881 a47eb4c Dawsh 2y ago
882 12b255b Dawsh 2y ago
883 a47eb4c Dawsh 2y ago def get_services(db: Session) -> List[Service]:
884 a47eb4c Dawsh 2y ago return db.query(Service).all()
885 a47eb4c Dawsh 2y ago
886 12b255b Dawsh 2y ago
887 f1ab16c Dawsh 2y ago def update_service(
888 f1ab16c Dawsh 2y ago db: Session, db_service: Service, modification: ServiceModify
889 f1ab16c Dawsh 2y ago ):
890 ee1351a Dawsh 2y ago if modification.name is not None:
891 ee1351a Dawsh 2y ago db_service.name = modification.name
892 ee1351a Dawsh 2y ago
893 db68a58 Dawsh 2y ago if modification.inbound_ids is not None:
894 01cdf2e Dawsh 2y ago db_service.inbounds = (
895 f1ab16c Dawsh 2y ago db.query(Inbound)
896 f1ab16c Dawsh 2y ago .filter(Inbound.id.in_(modification.inbound_ids))
897 f1ab16c Dawsh 2y ago .all()
898 01cdf2e Dawsh 2y ago )
899 ee1351a Dawsh 2y ago
900 ee1351a Dawsh 2y ago db.commit()
901 ee1351a Dawsh 2y ago db.refresh(db_service)
902 ee1351a Dawsh 2y ago return db_service
903 ee1351a Dawsh 2y ago
904 ee1351a Dawsh 2y ago
905 ee1351a Dawsh 2y ago def remove_service(db: Session, db_service: Service):
906 ee1351a Dawsh 2y ago db.delete(db_service)
907 ee1351a Dawsh 2y ago db.commit()
908 ee1351a Dawsh 2y ago return db_service
909 ee1351a Dawsh 2y ago
910 ee1351a Dawsh 2y ago
911 a47eb4c Dawsh 2y ago def get_node(db: Session, name: str):
912 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.name == name).first()
913 a47eb4c Dawsh 2y ago
914 a47eb4c Dawsh 2y ago
915 a47eb4c Dawsh 2y ago def get_node_by_id(db: Session, node_id: int):
916 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.id == node_id).first()
917 a47eb4c Dawsh 2y ago
918 a47eb4c Dawsh 2y ago
919 01cdf2e Dawsh 2y ago def get_nodes(
920 f1ab16c Dawsh 2y ago db: Session,
921 f1ab16c Dawsh 2y ago status: Optional[Union[NodeStatus, list]] = None,
922 f1ab16c Dawsh 2y ago enabled: bool = None,
923 01cdf2e Dawsh 2y ago ):
924 a47eb4c Dawsh 2y ago query = db.query(Node)
925 a47eb4c Dawsh 2y ago
926 a47eb4c Dawsh 2y ago if status:
927 a47eb4c Dawsh 2y ago if isinstance(status, list):
928 a47eb4c Dawsh 2y ago query = query.filter(Node.status.in_(status))
929 a47eb4c Dawsh 2y ago else:
930 a47eb4c Dawsh 2y ago query = query.filter(Node.status == status)
931 a47eb4c Dawsh 2y ago
932 a47eb4c Dawsh 2y ago if enabled:
933 a47eb4c Dawsh 2y ago query = query.filter(Node.status != NodeStatus.disabled)
934 a47eb4c Dawsh 2y ago
935 a47eb4c Dawsh 2y ago return query.all()
936 a47eb4c Dawsh 2y ago
937 a47eb4c Dawsh 2y ago
938 78a3e64 Dawsh 2y ago def get_node_usage(
939 78a3e64 Dawsh 2y ago db: Session, start: datetime, end: datetime, node: Node
940 78a3e64 Dawsh 2y ago ) -> TrafficUsageSeries:
941 78a3e64 Dawsh 2y ago usages = defaultdict(int)
942 a47eb4c Dawsh 2y ago
943 78a3e64 Dawsh 2y ago query = (
944 78a3e64 Dawsh 2y ago db.query(
945 78a3e64 Dawsh 2y ago NodeUserUsage.created_at, func.sum(NodeUserUsage.used_traffic)
946 78a3e64 Dawsh 2y ago )
947 78a3e64 Dawsh 2y ago .group_by(NodeUserUsage.created_at)
948 78a3e64 Dawsh 2y ago .filter(
949 78a3e64 Dawsh 2y ago and_(
950 78a3e64 Dawsh 2y ago NodeUserUsage.node_id == node.id,
951 78a3e64 Dawsh 2y ago NodeUserUsage.created_at >= start,
952 78a3e64 Dawsh 2y ago NodeUserUsage.created_at <= end,
953 78a3e64 Dawsh 2y ago )
954 78a3e64 Dawsh 2y ago )
955 a47eb4c Dawsh 2y ago )
956 78a3e64 Dawsh 2y ago
957 78a3e64 Dawsh 2y ago for created_at, used_traffic in query.all():
958 78a3e64 Dawsh 2y ago usages[created_at.replace(tzinfo=timezone.utc).timestamp()] += int(
959 78a3e64 Dawsh 2y ago used_traffic
960 a47eb4c Dawsh 2y ago )
961 a47eb4c Dawsh 2y ago
962 78a3e64 Dawsh 2y ago result = TrafficUsageSeries(usages=[], total=0)
963 78a3e64 Dawsh 2y ago current = start.astimezone(timezone.utc).replace(
964 78a3e64 Dawsh 2y ago minute=0, second=0, microsecond=0
965 78a3e64 Dawsh 2y ago )
966 a47eb4c Dawsh 2y ago
967 78a3e64 Dawsh 2y ago while current <= end.replace(tzinfo=timezone.utc):
968 78a3e64 Dawsh 2y ago usage = usages.get(current.timestamp()) or 0
969 b4ec790 Dawsh 2y ago result.usages.append((int(current.timestamp()), usage))
970 78a3e64 Dawsh 2y ago result.total += usage
971 78a3e64 Dawsh 2y ago current += timedelta(hours=1)
972 a47eb4c Dawsh 2y ago
973 78a3e64 Dawsh 2y ago return result
974 a47eb4c Dawsh 2y ago
975 a47eb4c Dawsh 2y ago
976 a47eb4c Dawsh 2y ago def create_node(db: Session, node: NodeCreate):
977 01cdf2e Dawsh 2y ago dbnode = Node(
978 01cdf2e Dawsh 2y ago name=node.name,
979 01cdf2e Dawsh 2y ago address=node.address,
980 01cdf2e Dawsh 2y ago port=node.port,
981 01cdf2e Dawsh 2y ago connection_backend=node.connection_backend,
982 01cdf2e Dawsh 2y ago )
983 a47eb4c Dawsh 2y ago
984 a47eb4c Dawsh 2y ago db.add(dbnode)
985 a47eb4c Dawsh 2y ago db.commit()
986 a47eb4c Dawsh 2y ago db.refresh(dbnode)
987 a47eb4c Dawsh 2y ago return dbnode
988 a47eb4c Dawsh 2y ago
989 a47eb4c Dawsh 2y ago
990 a47eb4c Dawsh 2y ago def remove_node(db: Session, dbnode: Node):
991 a47eb4c Dawsh 2y ago db.delete(dbnode)
992 a47eb4c Dawsh 2y ago db.commit()
993 a47eb4c Dawsh 2y ago return dbnode
994 a47eb4c Dawsh 2y ago
995 a47eb4c Dawsh 2y ago
996 a47eb4c Dawsh 2y ago def update_node(db: Session, dbnode: Node, modify: NodeModify):
997 a47eb4c Dawsh 2y ago if modify.name is not None:
998 a47eb4c Dawsh 2y ago dbnode.name = modify.name
999 a47eb4c Dawsh 2y ago
1000 a47eb4c Dawsh 2y ago if modify.address is not None:
1001 a47eb4c Dawsh 2y ago dbnode.address = modify.address
1002 a47eb4c Dawsh 2y ago
1003 a47eb4c Dawsh 2y ago if modify.port is not None:
1004 a47eb4c Dawsh 2y ago dbnode.port = modify.port
1005 a47eb4c Dawsh 2y ago
1006 a47eb4c Dawsh 2y ago if modify.status is NodeStatus.disabled:
1007 a47eb4c Dawsh 2y ago dbnode.status = modify.status
1008 a47eb4c Dawsh 2y ago dbnode.xray_version = None
1009 a47eb4c Dawsh 2y ago dbnode.message = None
1010 a47eb4c Dawsh 2y ago else:
1011 12b255b Dawsh 2y ago dbnode.status = NodeStatus.unhealthy
1012 a47eb4c Dawsh 2y ago
1013 a47eb4c Dawsh 2y ago if modify.usage_coefficient:
1014 a47eb4c Dawsh 2y ago dbnode.usage_coefficient = modify.usage_coefficient
1015 a47eb4c Dawsh 2y ago
1016 5584d4a Mrdn 2y ago if modify.connection_backend:
1017 5584d4a Mrdn 2y ago dbnode.connection_backend = modify.connection_backend
1018 5584d4a Mrdn 2y ago
1019 a47eb4c Dawsh 2y ago db.commit()
1020 a47eb4c Dawsh 2y ago db.refresh(dbnode)
1021 a47eb4c Dawsh 2y ago return dbnode
1022 a47eb4c Dawsh 2y ago
1023 a47eb4c Dawsh 2y ago
1024 01cdf2e Dawsh 2y ago def update_node_status(
1025 01cdf2e Dawsh 2y ago db: Session,
1026 01cdf2e Dawsh 2y ago node_id: int,
1027 01cdf2e Dawsh 2y ago status: NodeStatus,
1028 01cdf2e Dawsh 2y ago message: str = None,
1029 01cdf2e Dawsh 2y ago version: str = None,
1030 01cdf2e Dawsh 2y ago ):
1031 f6a32ea Dawsh 2y ago db_node = db.query(Node).where(Node.id == node_id).first()
1032 f6a32ea Dawsh 2y ago db_node.status = status
1033 f6a32ea Dawsh 2y ago if message:
1034 f6a32ea Dawsh 2y ago db_node.message = message
1035 f6a32ea Dawsh 2y ago db_node.last_status_change = datetime.utcnow()
1036 a47eb4c Dawsh 2y ago db.commit()