gitscope
/
← marzneshin / app/db/crud.py / blame
Files Commits 446 Branches 1 Tags 43 Markers 4 README License
https://dev.kronoz.win/r/marzneshin.git
app/db/crud.py · v0.1.2
Source
1 12b255b Dawsh 2y ago import json
2 34a7612 Dawsh 2y ago import secrets
3 a47eb4c Dawsh 2y ago from datetime import datetime
4 a47eb4c Dawsh 2y ago from enum import Enum
5 2de9d72 Dawsh 2y ago from typing import List, Optional, Tuple, Union
6 a47eb4c Dawsh 2y ago
7 a47eb4c Dawsh 2y ago from sqlalchemy import and_, delete
8 a47eb4c Dawsh 2y ago from sqlalchemy.orm import Session
9 a47eb4c Dawsh 2y ago
10 a47eb4c Dawsh 2y ago from app.db.models import (JWT, TLS, Admin, Node, NodeUsage, NodeUserUsage,
11 12b255b Dawsh 2y ago NotificationReminder,
12 a47eb4c Dawsh 2y ago InboundHost, Service,
13 2de9d72 Dawsh 2y ago Inbound, System, User,
14 12b255b Dawsh 2y ago UserUsageResetLogs)
15 a47eb4c Dawsh 2y ago from app.models.admin import AdminCreate, AdminModify, AdminPartialModify
16 a47eb4c Dawsh 2y ago from app.models.node import (NodeCreate, NodeModify, NodeStatus,
17 a47eb4c Dawsh 2y ago NodeUsageResponse)
18 a47eb4c Dawsh 2y ago from app.models.proxy import InboundHost as InboundHostModify
19 2de9d72 Dawsh 2y ago from app.models.service import Service as ServiceModify, ServiceCreate
20 a47eb4c Dawsh 2y ago from app.models.user import (ReminderType, UserCreate,
21 a47eb4c Dawsh 2y ago UserDataLimitResetStrategy, UserModify,
22 2de9d72 Dawsh 2y ago UserStatus, UserUsageResponse)
23 a47eb4c Dawsh 2y ago from app.utils.helpers import (calculate_expiration_days,
24 a47eb4c Dawsh 2y ago calculate_usage_percent)
25 a47eb4c Dawsh 2y ago from config import NOTIFY_DAYS_LEFT, NOTIFY_REACHED_USAGE_PERCENT
26 a47eb4c Dawsh 2y ago
27 a47eb4c Dawsh 2y ago
28 a47eb4c Dawsh 2y ago def add_default_hosts(db: Session, inbounds: List[Inbound]):
29 2de9d72 Dawsh 2y ago hosts = [InboundHost(remark="🚀 Marz ({USERNAME}) [{PROTOCOL} - {TRANSPORT}]", address="{SERVER_IP}", inbound=i) for
30 2de9d72 Dawsh 2y ago i in inbounds]
31 a47eb4c Dawsh 2y ago db.add_all(hosts)
32 a47eb4c Dawsh 2y ago db.commit()
33 a47eb4c Dawsh 2y ago
34 a47eb4c Dawsh 2y ago
35 a47eb4c Dawsh 2y ago def assure_node_inbounds(db: Session, inbounds: List[Inbound], node_id: int):
36 2de9d72 Dawsh 2y ago existing_inbounds = set(
37 2de9d72 Dawsh 2y ago (i.protocol, i.tag, i.config) for i in db.query(Inbound).filter(Inbound.node_id == node_id).all())
38 4d647dd Dawsh 2y ago inbounds_set = set((json.loads(i.config)["protocol"], i.tag, i.config) for i in list(inbounds))
39 12b255b Dawsh 2y ago new_inbounds = inbounds_set - existing_inbounds
40 12b255b Dawsh 2y ago deleted_inbounds = existing_inbounds - inbounds_set
41 a47eb4c Dawsh 2y ago # difference = [ i for i in existing_inbounds if (i.protocol, i.tag) not in existing_inbounds ]
42 12b255b Dawsh 2y ago if not new_inbounds and not deleted_inbounds:
43 a47eb4c Dawsh 2y ago return
44 12b255b Dawsh 2y ago """stmt = (
45 a47eb4c Dawsh 2y ago delete(Inbound)
46 12b255b Dawsh 2y ago .where(Inbound.node_id == node_id)
47 a47eb4c Dawsh 2y ago .execution_options(synchronize_session="fetch")
48 a47eb4c Dawsh 2y ago )
49 12b255b Dawsh 2y ago db.execute(stmt)"""
50 4d647dd Dawsh 2y ago new_inbounds = [Inbound(tag=i[1], protocol=i[0], config=i[2], node_id=node_id) for i in new_inbounds]
51 a47eb4c Dawsh 2y ago db.add_all(new_inbounds)
52 a47eb4c Dawsh 2y ago db.flush()
53 a47eb4c Dawsh 2y ago for i in new_inbounds:
54 2de9d72 Dawsh 2y ago db.refresh(i)
55 a47eb4c Dawsh 2y ago add_default_hosts(db, new_inbounds)
56 a47eb4c Dawsh 2y ago db.commit()
57 a47eb4c Dawsh 2y ago
58 a47eb4c Dawsh 2y ago
59 12b255b Dawsh 2y ago def get_node_users(db: Session, node_id: Optional[int], statuses: Optional[List[UserStatus]] = None):
60 12b255b Dawsh 2y ago query = db.query(User.id, User.username, User.key, Inbound).distinct(
61 2de9d72 Dawsh 2y ago ).join(Inbound.services).join(Service.users).filter(Inbound.node_id == node_id)
62 12b255b Dawsh 2y ago if isinstance(statuses, list):
63 a47eb4c Dawsh 2y ago query = query.filter(User.status.in_(statuses))
64 a47eb4c Dawsh 2y ago return query.all()
65 a47eb4c Dawsh 2y ago
66 12b255b Dawsh 2y ago
67 a47eb4c Dawsh 2y ago def get_user_hosts(db: Session, user_id: int):
68 2de9d72 Dawsh 2y ago return db.query(InboundHost).distinct().join(User.services).join(Service.inbounds).join(Inbound.hosts).filter(
69 2de9d72 Dawsh 2y ago User.id == user_id).all()
70 a47eb4c Dawsh 2y ago
71 12b255b Dawsh 2y ago
72 a47eb4c Dawsh 2y ago def get_inbound_hosts(db: Session, inbound_id: int) -> List[InboundHost]:
73 a47eb4c Dawsh 2y ago return db.query(InboundHost).filter(InboundHost.inbound_id == inbound_id).all()
74 a47eb4c Dawsh 2y ago
75 12b255b Dawsh 2y ago
76 a47eb4c Dawsh 2y ago def get_all_inbounds(db: Session):
77 a47eb4c Dawsh 2y ago return db.query(Inbound).all()
78 a47eb4c Dawsh 2y ago
79 a47eb4c Dawsh 2y ago
80 2de9d72 Dawsh 2y ago def get_inbound(db: Session, inbound_id: int) -> Inbound | None:
81 2de9d72 Dawsh 2y ago return db.query(Inbound).filter(Inbound.id == inbound_id).first()
82 a47eb4c Dawsh 2y ago
83 a47eb4c Dawsh 2y ago
84 2de9d72 Dawsh 2y ago def get_host(db: Session, host_id) -> InboundHost:
85 2de9d72 Dawsh 2y ago return db.query(InboundHost).filter(InboundHost.id == host_id).first()
86 a47eb4c Dawsh 2y ago
87 a47eb4c Dawsh 2y ago
88 a47eb4c Dawsh 2y ago def add_host(db: Session, inbound: Inbound, host: InboundHostModify):
89 2de9d72 Dawsh 2y ago host = InboundHost(
90 2de9d72 Dawsh 2y ago remark=host.remark,
91 2de9d72 Dawsh 2y ago address=host.address,
92 2de9d72 Dawsh 2y ago port=host.port,
93 2de9d72 Dawsh 2y ago path=host.path,
94 2de9d72 Dawsh 2y ago sni=host.sni,
95 2de9d72 Dawsh 2y ago host=host.host,
96 2de9d72 Dawsh 2y ago inbound=inbound,
97 2de9d72 Dawsh 2y ago security=host.security,
98 2de9d72 Dawsh 2y ago alpn=host.alpn,
99 2de9d72 Dawsh 2y ago fingerprint=host.fingerprint
100 a47eb4c Dawsh 2y ago )
101 2de9d72 Dawsh 2y ago inbound.hosts.append(host)
102 a47eb4c Dawsh 2y ago db.commit()
103 2de9d72 Dawsh 2y ago db.refresh(host)
104 2de9d72 Dawsh 2y ago return host
105 2de9d72 Dawsh 2y ago
106 2de9d72 Dawsh 2y ago
107 2de9d72 Dawsh 2y ago def update_host(db: Session, db_host: InboundHost, host: InboundHostModify):
108 2de9d72 Dawsh 2y ago db_host.remark = host.remark
109 2de9d72 Dawsh 2y ago db_host.address = host.address
110 2de9d72 Dawsh 2y ago db_host.port = host.port
111 2de9d72 Dawsh 2y ago db_host.path = host.path
112 2de9d72 Dawsh 2y ago db_host.sni = host.sni
113 2de9d72 Dawsh 2y ago db_host.host = host.host
114 2de9d72 Dawsh 2y ago db_host.security = host.security
115 2de9d72 Dawsh 2y ago db_host.alpn = host.alpn
116 2de9d72 Dawsh 2y ago db_host.fingerprint = host.fingerprint
117 2de9d72 Dawsh 2y ago db.commit()
118 2de9d72 Dawsh 2y ago db.refresh(db_host)
119 2de9d72 Dawsh 2y ago return db_host
120 a47eb4c Dawsh 2y ago
121 a47eb4c Dawsh 2y ago
122 a47eb4c Dawsh 2y ago def update_hosts(db: Session, inbound_tag: str, modified_hosts: List[InboundHostModify]):
123 a47eb4c Dawsh 2y ago inbound = get_or_create_inbound(db, inbound_tag)
124 a47eb4c Dawsh 2y ago inbound.hosts = [
125 a47eb4c Dawsh 2y ago InboundHost(
126 a47eb4c Dawsh 2y ago remark=host.remark,
127 a47eb4c Dawsh 2y ago address=host.address,
128 a47eb4c Dawsh 2y ago port=host.port,
129 a47eb4c Dawsh 2y ago path=host.path,
130 a47eb4c Dawsh 2y ago sni=host.sni,
131 a47eb4c Dawsh 2y ago host=host.host,
132 a47eb4c Dawsh 2y ago inbound=inbound,
133 a47eb4c Dawsh 2y ago security=host.security,
134 a47eb4c Dawsh 2y ago alpn=host.alpn,
135 a47eb4c Dawsh 2y ago fingerprint=host.fingerprint,
136 a47eb4c Dawsh 2y ago allowinsecure=host.allowinsecure,
137 a47eb4c Dawsh 2y ago is_disabled=host.is_disabled
138 a47eb4c Dawsh 2y ago ) for host in modified_hosts
139 a47eb4c Dawsh 2y ago ]
140 a47eb4c Dawsh 2y ago db.commit()
141 a47eb4c Dawsh 2y ago db.refresh(inbound)
142 a47eb4c Dawsh 2y ago return inbound.hosts
143 a47eb4c Dawsh 2y ago
144 a47eb4c Dawsh 2y ago
145 a47eb4c Dawsh 2y ago def get_user(db: Session, username: str):
146 a47eb4c Dawsh 2y ago return db.query(User).filter(User.username == username).first()
147 a47eb4c Dawsh 2y ago
148 a47eb4c Dawsh 2y ago
149 a47eb4c Dawsh 2y ago def get_user_by_id(db: Session, user_id: int):
150 a47eb4c Dawsh 2y ago return db.query(User).filter(User.id == user_id).first()
151 a47eb4c Dawsh 2y ago
152 a47eb4c Dawsh 2y ago
153 a47eb4c Dawsh 2y ago UsersSortingOptions = Enum('UsersSortingOptions', {
154 a47eb4c Dawsh 2y ago 'username': User.username.asc(),
155 a47eb4c Dawsh 2y ago 'used_traffic': User.used_traffic.asc(),
156 a47eb4c Dawsh 2y ago 'data_limit': User.data_limit.asc(),
157 a47eb4c Dawsh 2y ago 'expire': User.expire.asc(),
158 a47eb4c Dawsh 2y ago 'created_at': User.created_at.asc(),
159 a47eb4c Dawsh 2y ago '-username': User.username.desc(),
160 a47eb4c Dawsh 2y ago '-used_traffic': User.used_traffic.desc(),
161 a47eb4c Dawsh 2y ago '-data_limit': User.data_limit.desc(),
162 a47eb4c Dawsh 2y ago '-expire': User.expire.desc(),
163 a47eb4c Dawsh 2y ago '-created_at': User.created_at.desc(),
164 a47eb4c Dawsh 2y ago })
165 a47eb4c Dawsh 2y ago
166 a47eb4c Dawsh 2y ago
167 a47eb4c Dawsh 2y ago def get_users(db: Session,
168 a47eb4c Dawsh 2y ago offset: Optional[int] = None,
169 a47eb4c Dawsh 2y ago limit: Optional[int] = None,
170 a47eb4c Dawsh 2y ago usernames: Optional[List[str]] = None,
171 a47eb4c Dawsh 2y ago status: Optional[Union[UserStatus, list]] = None,
172 a47eb4c Dawsh 2y ago sort: Optional[List[UsersSortingOptions]] = None,
173 a47eb4c Dawsh 2y ago admin: Optional[Admin] = None,
174 a47eb4c Dawsh 2y ago reset_strategy: Optional[Union[UserDataLimitResetStrategy, list]] = None,
175 a47eb4c Dawsh 2y ago return_with_count: bool = False) -> Union[List[User], Tuple[List[User], int]]:
176 a47eb4c Dawsh 2y ago query = db.query(User)
177 a47eb4c Dawsh 2y ago
178 a47eb4c Dawsh 2y ago if usernames:
179 a47eb4c Dawsh 2y ago if len(usernames) == 1:
180 a47eb4c Dawsh 2y ago query = query.filter(User.username.ilike(f"%{usernames[0]}%"))
181 a47eb4c Dawsh 2y ago else:
182 a47eb4c Dawsh 2y ago query = query.filter(User.username.in_(usernames))
183 a47eb4c Dawsh 2y ago
184 a47eb4c Dawsh 2y ago if status:
185 a47eb4c Dawsh 2y ago if isinstance(status, list):
186 a47eb4c Dawsh 2y ago query = query.filter(User.status.in_(status))
187 a47eb4c Dawsh 2y ago else:
188 a47eb4c Dawsh 2y ago query = query.filter(User.status == status)
189 a47eb4c Dawsh 2y ago
190 a47eb4c Dawsh 2y ago if reset_strategy:
191 a47eb4c Dawsh 2y ago if isinstance(reset_strategy, list):
192 a47eb4c Dawsh 2y ago query = query.filter(User.data_limit_reset_strategy.in_(reset_strategy))
193 a47eb4c Dawsh 2y ago else:
194 a47eb4c Dawsh 2y ago query = query.filter(User.data_limit_reset_strategy == reset_strategy)
195 a47eb4c Dawsh 2y ago
196 a47eb4c Dawsh 2y ago if admin:
197 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
198 a47eb4c Dawsh 2y ago
199 a47eb4c Dawsh 2y ago # count it before applying limit and offset
200 a47eb4c Dawsh 2y ago if return_with_count:
201 a47eb4c Dawsh 2y ago count = query.count()
202 a47eb4c Dawsh 2y ago
203 a47eb4c Dawsh 2y ago if sort:
204 a47eb4c Dawsh 2y ago query = query.order_by(*(opt.value for opt in sort))
205 a47eb4c Dawsh 2y ago
206 a47eb4c Dawsh 2y ago if offset:
207 a47eb4c Dawsh 2y ago query = query.offset(offset)
208 a47eb4c Dawsh 2y ago if limit:
209 a47eb4c Dawsh 2y ago query = query.limit(limit)
210 a47eb4c Dawsh 2y ago
211 a47eb4c Dawsh 2y ago if return_with_count:
212 a47eb4c Dawsh 2y ago return query.all(), count
213 a47eb4c Dawsh 2y ago
214 a47eb4c Dawsh 2y ago return query.all()
215 a47eb4c Dawsh 2y ago
216 a47eb4c Dawsh 2y ago
217 a47eb4c Dawsh 2y ago def get_user_usages(db: Session, dbuser: User, start: datetime, end: datetime,
218 a47eb4c Dawsh 2y ago ) -> List[UserUsageResponse]:
219 eb74041 Dawsh 2y ago usages = dict()
220 a47eb4c Dawsh 2y ago
221 a47eb4c Dawsh 2y ago for node in db.query(Node).all():
222 a47eb4c Dawsh 2y ago usages[node.id] = UserUsageResponse(
223 a47eb4c Dawsh 2y ago node_id=node.id,
224 a47eb4c Dawsh 2y ago node_name=node.name,
225 a47eb4c Dawsh 2y ago used_traffic=0
226 a47eb4c Dawsh 2y ago )
227 a47eb4c Dawsh 2y ago
228 a47eb4c Dawsh 2y ago cond = and_(NodeUserUsage.user_id == dbuser.id,
229 a47eb4c Dawsh 2y ago NodeUserUsage.created_at >= start,
230 a47eb4c Dawsh 2y ago NodeUserUsage.created_at <= end)
231 a47eb4c Dawsh 2y ago
232 a47eb4c Dawsh 2y ago for v in db.query(NodeUserUsage).filter(cond):
233 a47eb4c Dawsh 2y ago try:
234 a47eb4c Dawsh 2y ago usages[v.node_id or 0].used_traffic += v.used_traffic
235 a47eb4c Dawsh 2y ago except KeyError:
236 a47eb4c Dawsh 2y ago pass
237 a47eb4c Dawsh 2y ago
238 a47eb4c Dawsh 2y ago return list(usages.values())
239 a47eb4c Dawsh 2y ago
240 a47eb4c Dawsh 2y ago
241 a47eb4c Dawsh 2y ago def get_users_count(db: Session, status: UserStatus = None, admin: Admin = None):
242 a47eb4c Dawsh 2y ago query = db.query(User.id)
243 a47eb4c Dawsh 2y ago if admin:
244 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
245 a47eb4c Dawsh 2y ago if status:
246 a47eb4c Dawsh 2y ago query = query.filter(User.status == status)
247 a47eb4c Dawsh 2y ago return query.count()
248 a47eb4c Dawsh 2y ago
249 a47eb4c Dawsh 2y ago
250 a47eb4c Dawsh 2y ago def create_user(db: Session, user: UserCreate, admin: Admin = None):
251 a47eb4c Dawsh 2y ago dbuser = User(
252 a47eb4c Dawsh 2y ago username=user.username,
253 a47eb4c Dawsh 2y ago # proxies=proxies,
254 a47eb4c Dawsh 2y ago key=user.key,
255 2de9d72 Dawsh 2y ago services=db.query(Service).filter(Service.id.in_(user.services)).all(), # user.services,
256 a47eb4c Dawsh 2y ago status=user.status,
257 a47eb4c Dawsh 2y ago data_limit=(user.data_limit or None),
258 a47eb4c Dawsh 2y ago expire=(user.expire or None),
259 a47eb4c Dawsh 2y ago admin=admin,
260 a47eb4c Dawsh 2y ago data_limit_reset_strategy=user.data_limit_reset_strategy,
261 a47eb4c Dawsh 2y ago note=user.note,
262 a47eb4c Dawsh 2y ago on_hold_expire_duration=(user.on_hold_expire_duration or None),
263 a47eb4c Dawsh 2y ago on_hold_timeout=(user.on_hold_timeout or None),
264 a47eb4c Dawsh 2y ago )
265 a47eb4c Dawsh 2y ago db.add(dbuser)
266 a47eb4c Dawsh 2y ago db.commit()
267 a47eb4c Dawsh 2y ago db.refresh(dbuser)
268 a47eb4c Dawsh 2y ago return dbuser
269 a47eb4c Dawsh 2y ago
270 a47eb4c Dawsh 2y ago
271 a47eb4c Dawsh 2y ago def remove_user(db: Session, dbuser: User):
272 a47eb4c Dawsh 2y ago db.delete(dbuser)
273 a47eb4c Dawsh 2y ago # db.query(User).filter_by(id=user_id).delete()
274 a47eb4c Dawsh 2y ago db.commit()
275 a47eb4c Dawsh 2y ago
276 a47eb4c Dawsh 2y ago
277 a47eb4c Dawsh 2y ago def update_user(db: Session, dbuser: User, modify: UserModify):
278 a47eb4c Dawsh 2y ago if modify.status is not None:
279 a47eb4c Dawsh 2y ago dbuser.status = modify.status
280 a47eb4c Dawsh 2y ago
281 a47eb4c Dawsh 2y ago if modify.data_limit is not None:
282 a47eb4c Dawsh 2y ago dbuser.data_limit = (modify.data_limit or None)
283 a47eb4c Dawsh 2y ago if dbuser.status not in (UserStatus.expired, UserStatus.disabled):
284 a47eb4c Dawsh 2y ago if not dbuser.data_limit or dbuser.used_traffic < dbuser.data_limit:
285 a47eb4c Dawsh 2y ago if dbuser.status != UserStatus.on_hold:
286 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.active
287 a47eb4c Dawsh 2y ago
288 a47eb4c Dawsh 2y ago if not dbuser.data_limit or (calculate_usage_percent(
289 a47eb4c Dawsh 2y ago dbuser.used_traffic, dbuser.data_limit) < NOTIFY_REACHED_USAGE_PERCENT):
290 a47eb4c Dawsh 2y ago delete_notification_reminder_by_type(db, dbuser.id, ReminderType.data_usage)
291 a47eb4c Dawsh 2y ago else:
292 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.limited
293 a47eb4c Dawsh 2y ago
294 a47eb4c Dawsh 2y ago if modify.expire is not None:
295 a47eb4c Dawsh 2y ago dbuser.expire = (modify.expire or None)
296 a47eb4c Dawsh 2y ago if dbuser.status in (UserStatus.active, UserStatus.expired):
297 a47eb4c Dawsh 2y ago if not dbuser.expire or dbuser.expire > datetime.utcnow().timestamp():
298 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.active
299 a47eb4c Dawsh 2y ago if not dbuser.expire or (calculate_expiration_days(
300 a47eb4c Dawsh 2y ago dbuser.expire) > NOTIFY_DAYS_LEFT):
301 a47eb4c Dawsh 2y ago delete_notification_reminder_by_type(db, dbuser.id, ReminderType.expiration_date)
302 a47eb4c Dawsh 2y ago else:
303 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.expired
304 a47eb4c Dawsh 2y ago
305 a47eb4c Dawsh 2y ago if modify.note is not None:
306 a47eb4c Dawsh 2y ago dbuser.note = modify.note or None
307 a47eb4c Dawsh 2y ago
308 a47eb4c Dawsh 2y ago if modify.data_limit_reset_strategy is not None:
309 a47eb4c Dawsh 2y ago dbuser.data_limit_reset_strategy = modify.data_limit_reset_strategy.value
310 a47eb4c Dawsh 2y ago
311 a47eb4c Dawsh 2y ago if modify.on_hold_timeout is not None:
312 a47eb4c Dawsh 2y ago dbuser.on_hold_timeout = modify.on_hold_timeout
313 a47eb4c Dawsh 2y ago
314 a47eb4c Dawsh 2y ago if modify.on_hold_expire_duration is not None:
315 a47eb4c Dawsh 2y ago dbuser.on_hold_expire_duration = modify.on_hold_expire_duration
316 2de9d72 Dawsh 2y ago
317 a47eb4c Dawsh 2y ago if modify.services is not None:
318 a47eb4c Dawsh 2y ago dbuser.services = db.query(Service).filter(Service.id.in_(modify.services)).all()
319 a47eb4c Dawsh 2y ago dbuser.edit_at = datetime.utcnow()
320 a47eb4c Dawsh 2y ago
321 a47eb4c Dawsh 2y ago db.commit()
322 a47eb4c Dawsh 2y ago db.refresh(dbuser)
323 a47eb4c Dawsh 2y ago return dbuser
324 a47eb4c Dawsh 2y ago
325 a47eb4c Dawsh 2y ago
326 a47eb4c Dawsh 2y ago def reset_user_data_usage(db: Session, dbuser: User):
327 a47eb4c Dawsh 2y ago usage_log = UserUsageResetLogs(
328 a47eb4c Dawsh 2y ago user=dbuser,
329 a47eb4c Dawsh 2y ago used_traffic_at_reset=dbuser.used_traffic,
330 a47eb4c Dawsh 2y ago )
331 a47eb4c Dawsh 2y ago db.add(usage_log)
332 a47eb4c Dawsh 2y ago
333 a47eb4c Dawsh 2y ago dbuser.used_traffic = 0
334 a47eb4c Dawsh 2y ago dbuser.node_usages.clear()
335 6da1b8a Dawsh 2y ago if dbuser.status == UserStatus.limited:
336 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.active.value
337 a47eb4c Dawsh 2y ago db.add(dbuser)
338 a47eb4c Dawsh 2y ago
339 a47eb4c Dawsh 2y ago db.commit()
340 a47eb4c Dawsh 2y ago db.refresh(dbuser)
341 a47eb4c Dawsh 2y ago return dbuser
342 a47eb4c Dawsh 2y ago
343 a47eb4c Dawsh 2y ago
344 a47eb4c Dawsh 2y ago def revoke_user_sub(db: Session, dbuser: User):
345 34a7612 Dawsh 2y ago dbuser.key = secrets.token_hex(16)
346 a47eb4c Dawsh 2y ago db.commit()
347 a47eb4c Dawsh 2y ago db.refresh(dbuser)
348 a47eb4c Dawsh 2y ago return dbuser
349 a47eb4c Dawsh 2y ago
350 a47eb4c Dawsh 2y ago
351 a47eb4c Dawsh 2y ago def update_user_sub(db: Session, dbuser: User, user_agent: str):
352 a47eb4c Dawsh 2y ago dbuser.sub_updated_at = datetime.utcnow()
353 a47eb4c Dawsh 2y ago dbuser.sub_last_user_agent = user_agent
354 a47eb4c Dawsh 2y ago
355 a47eb4c Dawsh 2y ago db.commit()
356 a47eb4c Dawsh 2y ago db.refresh(dbuser)
357 a47eb4c Dawsh 2y ago return dbuser
358 a47eb4c Dawsh 2y ago
359 a47eb4c Dawsh 2y ago
360 a47eb4c Dawsh 2y ago def reset_all_users_data_usage(db: Session, admin: Optional[Admin] = None):
361 a47eb4c Dawsh 2y ago query = db.query(User)
362 a47eb4c Dawsh 2y ago
363 a47eb4c Dawsh 2y ago if admin:
364 a47eb4c Dawsh 2y ago query = query.filter(User.admin == admin)
365 a47eb4c Dawsh 2y ago
366 a47eb4c Dawsh 2y ago for dbuser in query.all():
367 a47eb4c Dawsh 2y ago dbuser.used_traffic = 0
368 a47eb4c Dawsh 2y ago if dbuser.status not in [UserStatus.on_hold, UserStatus.expired, UserStatus.disabled]:
369 a47eb4c Dawsh 2y ago dbuser.status = UserStatus.active
370 a47eb4c Dawsh 2y ago dbuser.usage_logs.clear()
371 a47eb4c Dawsh 2y ago dbuser.node_usages.clear()
372 a47eb4c Dawsh 2y ago db.add(dbuser)
373 a47eb4c Dawsh 2y ago
374 a47eb4c Dawsh 2y ago db.commit()
375 a47eb4c Dawsh 2y ago
376 a47eb4c Dawsh 2y ago
377 a47eb4c Dawsh 2y ago def update_user_status(db: Session, dbuser: User, status: UserStatus):
378 a47eb4c Dawsh 2y ago dbuser.status = status
379 a47eb4c Dawsh 2y ago db.commit()
380 a47eb4c Dawsh 2y ago db.refresh(dbuser)
381 a47eb4c Dawsh 2y ago return dbuser
382 a47eb4c Dawsh 2y ago
383 a47eb4c Dawsh 2y ago
384 a47eb4c Dawsh 2y ago def set_owner(db: Session, dbuser: User, admin: Admin):
385 a47eb4c Dawsh 2y ago dbuser.admin = admin
386 a47eb4c Dawsh 2y ago db.commit()
387 a47eb4c Dawsh 2y ago db.refresh(dbuser)
388 a47eb4c Dawsh 2y ago return dbuser
389 a47eb4c Dawsh 2y ago
390 a47eb4c Dawsh 2y ago
391 a47eb4c Dawsh 2y ago def start_user_expire(db: Session, dbuser: User):
392 a47eb4c Dawsh 2y ago expire = int(datetime.utcnow().timestamp()) + dbuser.on_hold_expire_duration
393 a47eb4c Dawsh 2y ago dbuser.expire = expire
394 a47eb4c Dawsh 2y ago db.commit()
395 a47eb4c Dawsh 2y ago db.refresh(dbuser)
396 a47eb4c Dawsh 2y ago return dbuser
397 a47eb4c Dawsh 2y ago
398 a47eb4c Dawsh 2y ago
399 a47eb4c Dawsh 2y ago def get_system_usage(db: Session):
400 a47eb4c Dawsh 2y ago return db.query(System).first()
401 a47eb4c Dawsh 2y ago
402 a47eb4c Dawsh 2y ago
403 a47eb4c Dawsh 2y ago def get_jwt_secret_key(db: Session):
404 a47eb4c Dawsh 2y ago return db.query(JWT).first().secret_key
405 a47eb4c Dawsh 2y ago
406 a47eb4c Dawsh 2y ago
407 a47eb4c Dawsh 2y ago def get_tls_certificate(db: Session):
408 a47eb4c Dawsh 2y ago return db.query(TLS).first()
409 a47eb4c Dawsh 2y ago
410 a47eb4c Dawsh 2y ago
411 a47eb4c Dawsh 2y ago def get_admin(db: Session, username: str):
412 a47eb4c Dawsh 2y ago return db.query(Admin).filter(Admin.username == username).first()
413 a47eb4c Dawsh 2y ago
414 a47eb4c Dawsh 2y ago
415 a47eb4c Dawsh 2y ago def create_admin(db: Session, admin: AdminCreate):
416 a47eb4c Dawsh 2y ago dbadmin = Admin(
417 a47eb4c Dawsh 2y ago username=admin.username,
418 a47eb4c Dawsh 2y ago hashed_password=admin.hashed_password,
419 a47eb4c Dawsh 2y ago is_sudo=admin.is_sudo
420 a47eb4c Dawsh 2y ago )
421 a47eb4c Dawsh 2y ago db.add(dbadmin)
422 a47eb4c Dawsh 2y ago db.commit()
423 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
424 a47eb4c Dawsh 2y ago return dbadmin
425 a47eb4c Dawsh 2y ago
426 a47eb4c Dawsh 2y ago
427 a47eb4c Dawsh 2y ago def update_admin(db: Session, dbadmin: Admin, modified_admin: AdminModify):
428 a47eb4c Dawsh 2y ago dbadmin.is_sudo = modified_admin.is_sudo
429 a47eb4c Dawsh 2y ago if dbadmin.hashed_password != modified_admin.hashed_password:
430 a47eb4c Dawsh 2y ago dbadmin.hashed_password = modified_admin.hashed_password
431 a47eb4c Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
432 a47eb4c Dawsh 2y ago db.commit()
433 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
434 2de9d72 Dawsh 2y ago return dbadmin
435 a47eb4c Dawsh 2y ago
436 a47eb4c Dawsh 2y ago
437 a47eb4c Dawsh 2y ago def partial_update_admin(db: Session, dbadmin: Admin, modified_admin: AdminPartialModify):
438 a47eb4c Dawsh 2y ago if modified_admin.is_sudo is not None:
439 a47eb4c Dawsh 2y ago dbadmin.is_sudo = modified_admin.is_sudo
440 a47eb4c Dawsh 2y ago if modified_admin.password is not None and dbadmin.hashed_password != modified_admin.hashed_password:
441 a47eb4c Dawsh 2y ago dbadmin.hashed_password = modified_admin.hashed_password
442 a47eb4c Dawsh 2y ago dbadmin.password_reset_at = datetime.utcnow()
443 a47eb4c Dawsh 2y ago
444 a47eb4c Dawsh 2y ago db.commit()
445 a47eb4c Dawsh 2y ago db.refresh(dbadmin)
446 a47eb4c Dawsh 2y ago return dbadmin
447 a47eb4c Dawsh 2y ago
448 a47eb4c Dawsh 2y ago
449 a47eb4c Dawsh 2y ago def remove_admin(db: Session, dbadmin: Admin):
450 a47eb4c Dawsh 2y ago db.delete(dbadmin)
451 a47eb4c Dawsh 2y ago db.commit()
452 a47eb4c Dawsh 2y ago return dbadmin
453 a47eb4c Dawsh 2y ago
454 a47eb4c Dawsh 2y ago
455 a47eb4c Dawsh 2y ago def get_admins(db: Session,
456 a47eb4c Dawsh 2y ago offset: Optional[int] = None,
457 a47eb4c Dawsh 2y ago limit: Optional[int] = None,
458 a47eb4c Dawsh 2y ago username: Optional[str] = None):
459 a47eb4c Dawsh 2y ago query = db.query(Admin)
460 a47eb4c Dawsh 2y ago if username:
461 a47eb4c Dawsh 2y ago query = query.filter(Admin.username.ilike(f'%{username}%'))
462 a47eb4c Dawsh 2y ago if offset:
463 a47eb4c Dawsh 2y ago query = query.offset(offset)
464 a47eb4c Dawsh 2y ago if limit:
465 a47eb4c Dawsh 2y ago query = query.limit(limit)
466 a47eb4c Dawsh 2y ago return query.all()
467 a47eb4c Dawsh 2y ago
468 a47eb4c Dawsh 2y ago
469 ee1351a Dawsh 2y ago def create_service(db: Session, service: ServiceCreate) -> Service:
470 a47eb4c Dawsh 2y ago dbservice = Service(
471 a47eb4c Dawsh 2y ago name=service.name,
472 a47eb4c Dawsh 2y ago inbounds=db.query(Inbound).filter(Inbound.id.in_(service.inbounds)).all(),
473 a47eb4c Dawsh 2y ago users=db.query(User).filter(User.id.in_(service.users)).all()
474 2de9d72 Dawsh 2y ago )
475 a47eb4c Dawsh 2y ago db.add(dbservice)
476 a47eb4c Dawsh 2y ago db.commit()
477 a47eb4c Dawsh 2y ago db.refresh(dbservice)
478 a47eb4c Dawsh 2y ago return dbservice
479 a47eb4c Dawsh 2y ago
480 a47eb4c Dawsh 2y ago
481 ee1351a Dawsh 2y ago def get_service(db: Session, service_id: id) -> Service:
482 ee1351a Dawsh 2y ago return db.query(Service).filter(Service.id == service_id).first()
483 a47eb4c Dawsh 2y ago
484 12b255b Dawsh 2y ago
485 a47eb4c Dawsh 2y ago def get_services(db: Session) -> List[Service]:
486 a47eb4c Dawsh 2y ago return db.query(Service).all()
487 a47eb4c Dawsh 2y ago
488 12b255b Dawsh 2y ago
489 ee1351a Dawsh 2y ago def update_service(db: Session, db_service: Service, modification: ServiceModify):
490 ee1351a Dawsh 2y ago if modification.name is not None:
491 ee1351a Dawsh 2y ago db_service.name = modification.name
492 ee1351a Dawsh 2y ago
493 ee1351a Dawsh 2y ago if modification.inbounds is not None:
494 ee1351a Dawsh 2y ago db_service.inbounds = db.query(Inbound).filter(Inbound.id.in_(modification.inbounds)).all()
495 ee1351a Dawsh 2y ago
496 ee1351a Dawsh 2y ago db.commit()
497 ee1351a Dawsh 2y ago db.refresh(db_service)
498 ee1351a Dawsh 2y ago return db_service
499 ee1351a Dawsh 2y ago
500 ee1351a Dawsh 2y ago
501 ee1351a Dawsh 2y ago def remove_service(db: Session, db_service: Service):
502 ee1351a Dawsh 2y ago db.delete(db_service)
503 ee1351a Dawsh 2y ago db.commit()
504 ee1351a Dawsh 2y ago return db_service
505 ee1351a Dawsh 2y ago
506 ee1351a Dawsh 2y ago
507 a47eb4c Dawsh 2y ago def get_node(db: Session, name: str):
508 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.name == name).first()
509 a47eb4c Dawsh 2y ago
510 a47eb4c Dawsh 2y ago
511 a47eb4c Dawsh 2y ago def get_node_by_id(db: Session, node_id: int):
512 a47eb4c Dawsh 2y ago return db.query(Node).filter(Node.id == node_id).first()
513 a47eb4c Dawsh 2y ago
514 a47eb4c Dawsh 2y ago
515 a47eb4c Dawsh 2y ago def get_nodes(db: Session,
516 a47eb4c Dawsh 2y ago status: Optional[Union[NodeStatus, list]] = None,
517 a47eb4c Dawsh 2y ago enabled: bool = None):
518 a47eb4c Dawsh 2y ago query = db.query(Node)
519 a47eb4c Dawsh 2y ago
520 a47eb4c Dawsh 2y ago if status:
521 a47eb4c Dawsh 2y ago if isinstance(status, list):
522 a47eb4c Dawsh 2y ago query = query.filter(Node.status.in_(status))
523 a47eb4c Dawsh 2y ago else:
524 a47eb4c Dawsh 2y ago query = query.filter(Node.status == status)
525 a47eb4c Dawsh 2y ago
526 a47eb4c Dawsh 2y ago if enabled:
527 a47eb4c Dawsh 2y ago query = query.filter(Node.status != NodeStatus.disabled)
528 a47eb4c Dawsh 2y ago
529 a47eb4c Dawsh 2y ago return query.all()
530 a47eb4c Dawsh 2y ago
531 a47eb4c Dawsh 2y ago
532 a47eb4c Dawsh 2y ago def get_nodes_usage(db: Session, start: datetime, end: datetime) -> List[NodeUsageResponse]:
533 12b255b Dawsh 2y ago usages = dict()
534 a47eb4c Dawsh 2y ago
535 a47eb4c Dawsh 2y ago usages[0] = NodeUsageResponse( # Main Core
536 a47eb4c Dawsh 2y ago node_id=None,
537 a47eb4c Dawsh 2y ago node_name="Master",
538 a47eb4c Dawsh 2y ago uplink=0,
539 a47eb4c Dawsh 2y ago downlink=0
540 a47eb4c Dawsh 2y ago )
541 a47eb4c Dawsh 2y ago for node in db.query(Node).all():
542 a47eb4c Dawsh 2y ago usages[node.id] = NodeUsageResponse(
543 a47eb4c Dawsh 2y ago node_id=node.id,
544 a47eb4c Dawsh 2y ago node_name=node.name,
545 a47eb4c Dawsh 2y ago uplink=0,
546 a47eb4c Dawsh 2y ago downlink=0
547 a47eb4c Dawsh 2y ago )
548 a47eb4c Dawsh 2y ago
549 a47eb4c Dawsh 2y ago cond = and_(NodeUsage.created_at >= start, NodeUsage.created_at <= end)
550 a47eb4c Dawsh 2y ago
551 a47eb4c Dawsh 2y ago for v in db.query(NodeUsage).filter(cond):
552 a47eb4c Dawsh 2y ago try:
553 a47eb4c Dawsh 2y ago usages[v.node_id or 0].uplink += v.uplink
554 a47eb4c Dawsh 2y ago usages[v.node_id or 0].downlink += v.downlink
555 a47eb4c Dawsh 2y ago except KeyError:
556 a47eb4c Dawsh 2y ago pass
557 a47eb4c Dawsh 2y ago
558 a47eb4c Dawsh 2y ago return list(usages.values())
559 a47eb4c Dawsh 2y ago
560 a47eb4c Dawsh 2y ago
561 a47eb4c Dawsh 2y ago def create_node(db: Session, node: NodeCreate):
562 a47eb4c Dawsh 2y ago dbnode = Node(name=node.name,
563 a47eb4c Dawsh 2y ago address=node.address,
564 f6a32ea Dawsh 2y ago port=node.port,
565 f6a32ea Dawsh 2y ago connection_backend=node.connection_backend)
566 a47eb4c Dawsh 2y ago
567 a47eb4c Dawsh 2y ago db.add(dbnode)
568 a47eb4c Dawsh 2y ago db.commit()
569 a47eb4c Dawsh 2y ago db.refresh(dbnode)
570 a47eb4c Dawsh 2y ago return dbnode
571 a47eb4c Dawsh 2y ago
572 a47eb4c Dawsh 2y ago
573 a47eb4c Dawsh 2y ago def remove_node(db: Session, dbnode: Node):
574 a47eb4c Dawsh 2y ago db.delete(dbnode)
575 a47eb4c Dawsh 2y ago db.commit()
576 a47eb4c Dawsh 2y ago return dbnode
577 a47eb4c Dawsh 2y ago
578 a47eb4c Dawsh 2y ago
579 a47eb4c Dawsh 2y ago def update_node(db: Session, dbnode: Node, modify: NodeModify):
580 a47eb4c Dawsh 2y ago if modify.name is not None:
581 a47eb4c Dawsh 2y ago dbnode.name = modify.name
582 a47eb4c Dawsh 2y ago
583 a47eb4c Dawsh 2y ago if modify.address is not None:
584 a47eb4c Dawsh 2y ago dbnode.address = modify.address
585 a47eb4c Dawsh 2y ago
586 a47eb4c Dawsh 2y ago if modify.port is not None:
587 a47eb4c Dawsh 2y ago dbnode.port = modify.port
588 a47eb4c Dawsh 2y ago
589 a47eb4c Dawsh 2y ago if modify.status is NodeStatus.disabled:
590 a47eb4c Dawsh 2y ago dbnode.status = modify.status
591 a47eb4c Dawsh 2y ago dbnode.xray_version = None
592 a47eb4c Dawsh 2y ago dbnode.message = None
593 a47eb4c Dawsh 2y ago else:
594 12b255b Dawsh 2y ago dbnode.status = NodeStatus.unhealthy
595 a47eb4c Dawsh 2y ago
596 a47eb4c Dawsh 2y ago if modify.usage_coefficient:
597 a47eb4c Dawsh 2y ago dbnode.usage_coefficient = modify.usage_coefficient
598 a47eb4c Dawsh 2y ago
599 a47eb4c Dawsh 2y ago db.commit()
600 a47eb4c Dawsh 2y ago db.refresh(dbnode)
601 a47eb4c Dawsh 2y ago return dbnode
602 a47eb4c Dawsh 2y ago
603 a47eb4c Dawsh 2y ago
604 f6a32ea Dawsh 2y ago def update_node_status(db: Session, node_id: int, status: NodeStatus, message: str = None, version: str = None):
605 f6a32ea Dawsh 2y ago db_node = db.query(Node).where(Node.id == node_id).first()
606 f6a32ea Dawsh 2y ago db_node.status = status
607 f6a32ea Dawsh 2y ago if message:
608 f6a32ea Dawsh 2y ago db_node.message = message
609 f6a32ea Dawsh 2y ago db_node.last_status_change = datetime.utcnow()
610 a47eb4c Dawsh 2y ago db.commit()
611 a47eb4c Dawsh 2y ago
612 a47eb4c Dawsh 2y ago
613 a47eb4c Dawsh 2y ago def create_notification_reminder(
614 a47eb4c Dawsh 2y ago db: Session, reminder_type: ReminderType, expires_at: datetime, user_id: int) -> NotificationReminder:
615 a47eb4c Dawsh 2y ago reminder = NotificationReminder(type=reminder_type, expires_at=expires_at, user_id=user_id)
616 a47eb4c Dawsh 2y ago db.add(reminder)
617 a47eb4c Dawsh 2y ago db.commit()
618 a47eb4c Dawsh 2y ago db.refresh(reminder)
619 a47eb4c Dawsh 2y ago return reminder
620 a47eb4c Dawsh 2y ago
621 a47eb4c Dawsh 2y ago
622 a47eb4c Dawsh 2y ago def get_notification_reminder(
623 a47eb4c Dawsh 2y ago db: Session, user_id: int, reminder_type: ReminderType,
624 12b255b Dawsh 2y ago ) -> Optional[NotificationReminder]:
625 a47eb4c Dawsh 2y ago reminder = db.query(NotificationReminder).filter(
626 a47eb4c Dawsh 2y ago NotificationReminder.user_id == user_id).filter(
627 a47eb4c Dawsh 2y ago NotificationReminder.type == reminder_type).first()
628 a47eb4c Dawsh 2y ago if reminder is None:
629 a47eb4c Dawsh 2y ago return
630 a47eb4c Dawsh 2y ago if reminder.expires_at and reminder.expires_at < datetime.utcnow():
631 a47eb4c Dawsh 2y ago db.delete(reminder)
632 a47eb4c Dawsh 2y ago db.commit()
633 a47eb4c Dawsh 2y ago return
634 a47eb4c Dawsh 2y ago return reminder
635 a47eb4c Dawsh 2y ago
636 a47eb4c Dawsh 2y ago
637 a47eb4c Dawsh 2y ago def delete_notification_reminder_by_type(db: Session, user_id: int, reminder_type: ReminderType) -> None:
638 a47eb4c Dawsh 2y ago """Deletes notification reminder filtered by user_id and type if exists"""
639 a47eb4c Dawsh 2y ago stmt = delete(NotificationReminder).where(
640 a47eb4c Dawsh 2y ago NotificationReminder.user_id == user_id,
641 a47eb4c Dawsh 2y ago NotificationReminder.type == reminder_type,
642 a47eb4c Dawsh 2y ago )
643 a47eb4c Dawsh 2y ago db.execute(stmt)
644 a47eb4c Dawsh 2y ago db.commit()
645 a47eb4c Dawsh 2y ago return
646 a47eb4c Dawsh 2y ago
647 a47eb4c Dawsh 2y ago
648 a47eb4c Dawsh 2y ago def delete_notification_reminder(db: Session, dbreminder: NotificationReminder) -> None:
649 a47eb4c Dawsh 2y ago db.delete(dbreminder)
650 a47eb4c Dawsh 2y ago db.commit()
651 a47eb4c Dawsh 2y ago return