Skip to content
This repository was archived by the owner on Jul 8, 2026. It is now read-only.

Commit 287ea9f

Browse files
ReturnFIReturnFI
committed
refactor(scripts): migrate IP limit state from JSON to MongoDB
1 parent fc55a39 commit 287ea9f

1 file changed

Lines changed: 69 additions & 151 deletions

File tree

‎core/scripts/hysteria2/limit.sh‎

Lines changed: 69 additions & 151 deletions
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,22 @@ source /etc/hysteria/core/scripts/path.sh
44

55
# --- Configuration ---
66
SERVICE_NAME="hysteria-ip-limit.service"
7+
DB_NAME="blitz_panel"
8+
CONNECTIONS_COLLECTION="active_connections"
79

810
# Load configurations from .configs.env
911
if [ -f "$CONFIG_ENV" ]; then
1012
source "$CONFIG_ENV"
11-
BLOCK_DURATION="${BLOCK_DURATION:-60}" # Default to 60 seconds if not set
12-
MAX_IPS="${MAX_IPS:-1}" # Default to 1 IP if not set
13+
BLOCK_DURATION="${BLOCK_DURATION:-60}" # Default to 60 seconds
14+
MAX_IPS="${MAX_IPS:-1}" # Default to 1 IP
1315

1416
grep -q "^BLOCK_DURATION=" "$CONFIG_ENV" || echo -e "\nBLOCK_DURATION=$BLOCK_DURATION" >> "$CONFIG_ENV"
15-
1617
grep -q "^MAX_IPS=" "$CONFIG_ENV" || echo "MAX_IPS=$MAX_IPS" >> "$CONFIG_ENV"
1718
else
1819
echo -e "BLOCK_DURATION=240\nMAX_IPS=5" > "$CONFIG_ENV"
1920
fi
2021

2122
# --- Ensure files exist ---
22-
[ ! -f "$CONNECTIONS_FILE" ] && echo "{}" > "$CONNECTIONS_FILE"
2323
[ ! -f "$BLOCK_LIST" ] && touch "$BLOCK_LIST"
2424

2525
# --- Logging function ---
@@ -29,69 +29,36 @@ log_message() {
2929
echo "[$(date +"%Y-%m-%d %H:%M:%S")] [$level] $message"
3030
}
3131

32-
# --- Function to update the JSON file with new connection data ---
33-
update_json() {
32+
# --- Add an IP to the database for a user ---
33+
add_ip_to_db() {
3434
local username="$1"
3535
local ip_address="$2"
36-
37-
if command -v jq &>/dev/null; then
38-
temp_file=$(mktemp)
39-
jq --arg user "$username" --arg ip "$ip_address" \
40-
'.[$user] += [$ip] | .[$user] |= unique' "$CONNECTIONS_FILE" > "$temp_file"
41-
mv "$temp_file" "$CONNECTIONS_FILE"
42-
else
43-
if grep -q "\"$username\"" "$CONNECTIONS_FILE"; then
44-
# Add IP to existing username (if it doesn't exist)
45-
if ! grep -q "\"$username\".*\"$ip_address\"" "$CONNECTIONS_FILE"; then
46-
sed -i -E "s/(\"$username\":\s*\[)([^\]]*)/\1\2,\"$ip_address\"/" "$CONNECTIONS_FILE"
47-
fi
48-
else
49-
# Add new username with IP
50-
sed -i -E "s/\{(.*)\}/{\1,\"$username\":[\"$ip_address\"]}/" "$CONNECTIONS_FILE"
51-
fi
52-
fi
53-
54-
log_message "INFO" "Updated JSON: Added $ip_address for user $username"
36+
37+
mongosh "$DB_NAME" --quiet --eval "
38+
db.getCollection('$CONNECTIONS_COLLECTION').updateOne(
39+
{ _id: '$username' },
40+
{ \$addToSet: { ips: '$ip_address' } },
41+
{ upsert: true }
42+
);
43+
"
44+
log_message "INFO" "DB Update: Added $ip_address for user $username"
5545
}
5646

57-
# --- Function to remove an IP from the JSON when client disconnects ---
58-
remove_ip() {
47+
# --- Remove an IP from the database for a user ---
48+
remove_ip_from_db() {
5949
local username="$1"
6050
local ip_address="$2"
61-
62-
if [ ! -f "$CONNECTIONS_FILE" ]; then
63-
log_message "ERROR" "JSON file does not exist"
64-
return
65-
fi
66-
67-
if grep -q "\"$username\"" "$CONNECTIONS_FILE"; then
68-
if command -v jq &>/dev/null; then
69-
temp_file=$(mktemp)
70-
jq --arg user "$username" --arg ip "$ip_address" \
71-
'.[$user] = (.[$user] | map(select(. != $ip)))' "$CONNECTIONS_FILE" > "$temp_file"
72-
mv "$temp_file" "$CONNECTIONS_FILE"
73-
74-
# Check if the user's IP list is now empty and remove the user if so
75-
temp_file_check=$(mktemp)
76-
jq --arg user "$username" 'if .[$user] | length == 0 then del(.[$user]) else . end' "$CONNECTIONS_FILE" > "$temp_file_check"
77-
mv "$temp_file_check" "$CONNECTIONS_FILE"
78-
79-
else
80-
# Basic sed replacement (not as reliable as jq)
81-
sed -i -E "s/\"$ip_address\"(,|\])|\1\"$ip_address\"/\1/g" "$CONNECTIONS_FILE"
82-
sed -i -E "s/,\s*\]/\]/g" "$CONNECTIONS_FILE"
83-
sed -i -E "s/\[\s*,/\[/g" "$CONNECTIONS_FILE"
84-
85-
# VERY Basic check if user's IP list is empty and remove the user if so (less reliable)
86-
if grep -q "\"$username\":\s*\[\s*\]" "$CONNECTIONS_FILE"; then
87-
sed -i "/\"$username\":\s*\[\s*\][,\s]*/d" "$CONNECTIONS_FILE"
88-
sed -i "s/,\s*\}$/\n}/" "$CONNECTIONS_FILE" # Remove trailing comma if it exists after user deletion
89-
fi
90-
fi
91-
log_message "INFO" "Updated JSON: Removed $ip_address for user $username"
92-
else
93-
log_message "WARN" "User $username not found in JSON"
94-
fi
51+
52+
mongosh "$DB_NAME" --quiet --eval "
53+
db.getCollection('$CONNECTIONS_COLLECTION').updateOne(
54+
{ _id: '$username' },
55+
{ \$pull: { ips: '$ip_address' } }
56+
);
57+
db.getCollection('$CONNECTIONS_COLLECTION').deleteMany(
58+
{ _id: '$username', ips: { \$size: 0 } }
59+
);
60+
"
61+
log_message "INFO" "DB Update: Removed $ip_address for user $username"
9562
}
9663

9764
# --- Block an IP using iptables and track it ---
@@ -100,60 +67,45 @@ block_ip() {
10067
local username="$2"
10168
local unblock_time=$(( $(date +%s) + BLOCK_DURATION ))
10269

103-
# Skip if already blocked
10470
if iptables -C INPUT -s "$ip_address" -j DROP 2>/dev/null; then
10571
log_message "INFO" "IP $ip_address is already blocked"
10672
return
10773
fi
10874

109-
# Add to iptables
11075
iptables -I INPUT -s "$ip_address" -j DROP
111-
112-
# Add to block list with expiration time
11376
echo "$ip_address,$username,$unblock_time" >> "$BLOCK_LIST"
114-
11577
log_message "WARN" "Blocked IP $ip_address for user $username for $BLOCK_DURATION seconds"
11678
}
11779

11880
# --- Explicitly unblock an IP using iptables ---
11981
unblock_ip() {
12082
local ip_address="$1"
12183

122-
# Remove from iptables if exists
12384
if iptables -C INPUT -s "$ip_address" -j DROP 2>/dev/null; then
12485
iptables -D INPUT -s "$ip_address" -j DROP
12586
log_message "INFO" "Unblocked IP $ip_address"
12687
fi
127-
128-
# Remove from block list
12988
sed -i "/$ip_address,/d" "$BLOCK_LIST"
13089
}
13190

13291
# --- Block all IPs for a user ---
13392
block_all_user_ips() {
13493
local username="$1"
135-
local ips=()
136-
137-
# Get all IPs for this user
138-
if command -v jq &>/dev/null; then
139-
readarray -t ips < <(jq -r --arg user "$username" '.[$user][]' "$CONNECTIONS_FILE" 2>/dev/null)
140-
else
141-
# Basic extraction without jq (less reliable)
142-
ip_list=$(grep -oP "\"$username\":\s*\[\K[^\]]*" "$CONNECTIONS_FILE")
143-
IFS=',' read -ra ip_entries <<< "$ip_list"
144-
for entry in "${ip_entries[@]}"; do
145-
# Extract IP from the JSON array entry
146-
ip=$(echo "$entry" | grep -oP '".*"' | tr -d '"' | tr -d '[:space:]')
147-
if [[ -n "$ip" ]]; then
148-
ips+=("$ip")
149-
fi
150-
done
94+
95+
local ips_json
96+
ips_json=$(mongosh "$DB_NAME" --quiet --eval "
97+
JSON.stringify(db.getCollection('$CONNECTIONS_COLLECTION').findOne({_id: '$username'}, {_id: 0, ips: 1}))
98+
")
99+
100+
if [[ -z "$ips_json" || "$ips_json" == "null" ]]; then
101+
log_message "INFO" "No IPs to block for user $username"
102+
return
151103
fi
152-
153-
# Block all IPs for this user
104+
105+
local ips
106+
readarray -t ips < <(echo "$ips_json" | jq -r '.ips[]')
107+
154108
for ip in "${ips[@]}"; do
155-
ip=${ip//\"/} # Remove quotes
156-
ip=$(echo "$ip" | tr -d '[:space:]') # Remove whitespace
157109
if [[ -n "$ip" ]]; then
158110
block_ip "$ip" "$username"
159111
fi
@@ -167,7 +119,6 @@ check_expired_blocks() {
167119
local current_time=$(date +%s)
168120
local ip username expiry
169121

170-
# Check each line in the block list
171122
while IFS=, read -r ip username expiry || [ -n "$ip" ]; do
172123
if [[ -n "$ip" && -n "$expiry" ]]; then
173124
if (( current_time >= expiry )); then
@@ -181,44 +132,22 @@ check_expired_blocks() {
181132
# --- Check if a user has exceeded the IP limit ---
182133
check_ip_limit() {
183134
local username="$1"
184-
local ips=()
185-
186-
local is_unlimited="false"
187-
if [ -f "$USERS_FILE" ]; then
188-
if command -v jq &>/dev/null; then
189-
is_unlimited=$(jq -r --arg user "$username" '.[$user].unlimited_user // "false"' "$USERS_FILE" 2>/dev/null)
190-
else
191-
if grep -q "\"$username\"" "$USERS_FILE" && \
192-
grep -A 5 "\"$username\"" "$USERS_FILE" | grep -q '"unlimited_user": true'; then
193-
is_unlimited="true"
194-
fi
195-
fi
196-
fi
197135

198-
if [ "$is_unlimited" = "true" ]; then
136+
local is_unlimited
137+
is_unlimited=$(mongosh "$DB_NAME" --quiet --eval "
138+
db.users.findOne({_id: '$username'}, {_id: 0, unlimited_user: 1})?.unlimited_user || false;
139+
")
140+
141+
if [ "$is_unlimited" == "true" ]; then
199142
log_message "INFO" "User $username is exempt from IP limit. Skipping check."
200143
return
201144
fi
145+
146+
local ip_count
147+
ip_count=$(mongosh "$DB_NAME" --quiet --eval "
148+
db.getCollection('$CONNECTIONS_COLLECTION').findOne({_id: '$username'})?.ips?.length || 0;
149+
")
202150

203-
# Get all IPs for this user
204-
if command -v jq &>/dev/null; then
205-
readarray -t ips < <(jq -r --arg user "$username" '.[$user][]' "$CONNECTIONS_FILE" 2>/dev/null)
206-
else
207-
# Basic extraction without jq (less reliable)
208-
ip_list=$(grep -oP "\"$username\":\s*\[\K[^\]]*" "$CONNECTIONS_FILE")
209-
IFS=',' read -ra ip_entries <<< "$ip_list"
210-
for entry in "${ip_entries[@]}"; do
211-
# Extract IP from the JSON array entry
212-
ip=$(echo "$entry" | grep -oP '".*"' | tr -d '"' | tr -d '[:space:]')
213-
if [[ -n "$ip" ]]; then
214-
ips+=("$ip")
215-
fi
216-
done
217-
fi
218-
219-
ip_count=${#ips[@]}
220-
221-
# If the user has more IPs than allowed, block ALL their IPs
222151
if (( ip_count > MAX_IPS )); then
223152
log_message "WARN" "User $username has $ip_count IPs (max: $MAX_IPS) - blocking all IPs"
224153
block_all_user_ips "$username"
@@ -228,29 +157,25 @@ check_ip_limit() {
228157
# --- Parse log lines for connections and disconnections ---
229158
parse_log_line() {
230159
local log_line="$1"
231-
local ip_address=""
232-
local username=""
160+
local ip_address
161+
local username
233162

234-
# Extract IP address and username
235163
ip_address=$(echo "$log_line" | grep -oP '"addr": "([^:]+)' | cut -d'"' -f4)
236164
username=$(echo "$log_line" | grep -oP '"id": "([^">]+)' | cut -d'"' -f4)
237165

238166
if [[ -n "$username" && -n "$ip_address" ]]; then
239167
if echo "$log_line" | grep -q "client connected"; then
240-
# Check if this IP is in the block list
241168
if grep -q "^$ip_address," "$BLOCK_LIST"; then
242169
log_message "WARN" "Rejected connection from blocked IP $ip_address for user $username"
243-
# Make sure the IP is still blocked in iptables
244170
if ! iptables -C INPUT -s "$ip_address" -j DROP 2>/dev/null; then
245171
iptables -I INPUT -s "$ip_address" -j DROP
246172
fi
247173
else
248-
update_json "$username" "$ip_address"
174+
add_ip_to_db "$username" "$ip_address"
249175
check_ip_limit "$username"
250176
fi
251177
elif echo "$log_line" | grep -q "client disconnected"; then
252-
remove_ip "$username" "$ip_address"
253-
# Note: We don't unblock on disconnect - only on block expiration
178+
remove_ip_from_db "$username" "$ip_address"
254179
fi
255180
fi
256181
}
@@ -259,9 +184,9 @@ parse_log_line() {
259184
install_service() {
260185
cat <<EOF > /etc/systemd/system/${SERVICE_NAME}
261186
[Unit]
262-
Description=Hysteria2 IP Limiter
263-
After=network.target hysteria-server.service
264-
Requires=hysteria-server.service
187+
Description=Hysteria2 IP Limiter (MongoDB version)
188+
After=network.target hysteria-server.service mongod.service
189+
Requires=hysteria-server.service mongod.service
265190
266191
[Service]
267192
Type=simple
@@ -319,19 +244,20 @@ change_config() {
319244
fi
320245
}
321246

322-
# --- Check if running as root ---
247+
# --- Startup Checks ---
323248
if [[ $EUID -ne 0 ]]; then
324-
echo "Error: This script must be run as root for iptables functionality."
249+
echo "Error: This script must be run as root."
250+
exit 1
251+
fi
252+
if ! command -v mongosh &>/dev/null; then
253+
log_message "ERROR" "'mongosh' is not installed or not in PATH. This script requires the MongoDB Shell."
325254
exit 1
326255
fi
327-
328-
# --- Check for jq and warn if not available ---
329256
if ! command -v jq &>/dev/null; then
330-
log_message "WARN" "'jq' is not installed. JSON handling may be less reliable."
331-
log_message "WARN" "Consider installing jq with: apt install jq (for Debian/Ubuntu)"
257+
log_message "WARN" "'jq' is not installed. JSON parsing for blocking might fail."
332258
fi
333259

334-
# --- Command execution based on arguments ---
260+
# --- Command execution ---
335261
case "$1" in
336262
start)
337263
install_service
@@ -343,32 +269,24 @@ case "$1" in
343269
change_config "$2" "$3"
344270
;;
345271
run)
346-
log_message "INFO" "Monitoring Hysteria server connections. Max IPs per user: $MAX_IPS"
347-
log_message "INFO" "Block duration: $BLOCK_DURATION seconds"
348-
log_message "INFO" "Connection data saved to: $CONNECTIONS_FILE"
349-
log_message "INFO" "Press Ctrl+C to exit"
272+
log_message "INFO" "Monitoring Hysteria connections. Max IPs: $MAX_IPS, Block Duration: $BLOCK_DURATION s"
350273
log_message "INFO" "--------------------------------------------------------"
351274

352-
# Background process to check for expired blocks every 10 seconds
353275
(
354276
while true; do
355277
check_expired_blocks
356278
sleep 10
357279
done
358280
) &
359281
CHECKER_PID=$!
360-
361-
# Cleanup function
282+
362283
cleanup() {
363284
log_message "INFO" "Stopping IP limiter..."
364285
kill $CHECKER_PID 2>/dev/null
365286
exit 0
366287
}
367-
368-
# Set trap for cleanup
369288
trap cleanup SIGINT SIGTERM
370289

371-
# Monitor log for connections and disconnections
372290
journalctl -u hysteria-server.service -f | while read -r line; do
373291
if echo "$line" | grep -q "client connected\|client disconnected"; then
374292
parse_log_line "$line"

0 commit comments

Comments
 (0)