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