AstraOS commited on
Commit
ecb74c0
Β·
verified Β·
1 Parent(s): 360c1c2

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +463 -0
app.py ADDED
@@ -0,0 +1,463 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import telebot
2
+ import subprocess
3
+ import threading
4
+ import time
5
+ import re
6
+ import html
7
+ from collections import deque
8
+ import os
9
+ import re
10
+ import signal
11
+ import psutil
12
+ from fastapi import FastAPI
13
+
14
+
15
+ # Initialize bot with your token
16
+ TOKEN = os.environ["BOT_TOKEN"]
17
+ bot = telebot.TeleBot(TOKEN)
18
+
19
+ # Store authorized chat IDs
20
+ AUTHORIZED_USERS = list(map(int, re.findall(r"\d+", str(os.getenv("CHAT_IDS", "")))))
21
+
22
+ # Store active command processes
23
+ active_processes = {}
24
+ # Store logs history (for each chat_id)
25
+ logs_history = {}
26
+ MESSAGE_LIMIT = 4000
27
+
28
+ # Maximum number of log lines to keep in history
29
+ MAX_LOG_HISTORY = 1000
30
+ # Maximum number of command history entries
31
+ MAX_COMMAND_HISTORY = 20
32
+
33
+ # Store command history
34
+ command_history = {}
35
+
36
+ def is_authorized(chat_id):
37
+ """Check if user is authorized"""
38
+ return chat_id in AUTHORIZED_USERS
39
+
40
+ def initialize_logs_history(chat_id):
41
+ """Initialize logs history for a chat"""
42
+ if chat_id not in logs_history:
43
+ logs_history[chat_id] = deque(maxlen=MAX_LOG_HISTORY)
44
+
45
+ def add_to_logs_history(chat_id, log_line):
46
+ """Add log line to history"""
47
+ initialize_logs_history(chat_id)
48
+ logs_history[chat_id].append(log_line)
49
+
50
+ def initialize_command_history(chat_id):
51
+ """Initialize command history for a chat"""
52
+ if chat_id not in command_history:
53
+ command_history[chat_id] = deque(maxlen=MAX_COMMAND_HISTORY)
54
+
55
+ def add_to_command_history(chat_id, command):
56
+ """Add a command to the history"""
57
+ initialize_command_history(chat_id)
58
+ command_history[chat_id].append(command)
59
+
60
+ def get_command_history(chat_id):
61
+ """Get the command history for a chat"""
62
+ if chat_id in command_history:
63
+ return list(command_history[chat_id])
64
+ return []
65
+
66
+ def get_latest_logs(chat_id):
67
+ """Get the latest logs for a chat"""
68
+ if chat_id in logs_history:
69
+ return "\n".join(logs_history[chat_id])
70
+ return ""
71
+
72
+ def sanitize_log(text):
73
+ """Sanitize text for HTML formatting"""
74
+ return html.escape(text)
75
+
76
+ def split_logs(logs, latest_only=False):
77
+ """Split logs into chunks respecting MESSAGE_LIMIT.
78
+ If `latest_only` is True, only return the last chunk.
79
+ """
80
+ chunks = []
81
+ current_chunk = ""
82
+
83
+ for line in logs.split('\n'):
84
+ if len(current_chunk) + len(line) + 1 > MESSAGE_LIMIT:
85
+ chunks.append(current_chunk)
86
+ current_chunk = line + '\n'
87
+ else:
88
+ current_chunk += line + '\n'
89
+
90
+ if current_chunk:
91
+ chunks.append(current_chunk)
92
+
93
+ # If latest_only is True, return only the last chunk
94
+ if latest_only and chunks:
95
+ return [chunks[-1]]
96
+ return chunks
97
+
98
+
99
+ def create_log_keyboard(is_live=True):
100
+ """Create inline keyboard for log control"""
101
+ keyboard = telebot.types.InlineKeyboardMarkup(row_width=2)
102
+ if is_live:
103
+ keyboard.add(
104
+ telebot.types.InlineKeyboardButton("⏸️ Pause", callback_data="stop_logs"),
105
+ telebot.types.InlineKeyboardButton("▢️ Resume", callback_data="resume_logs"),
106
+ telebot.types.InlineKeyboardButton("πŸ”„ Refresh", callback_data="refresh_logs"),
107
+ telebot.types.InlineKeyboardButton("❌ Clear", callback_data="clear_logs")
108
+ )
109
+ else:
110
+ keyboard.add(
111
+ telebot.types.InlineKeyboardButton("❌ Clear log history", callback_data="clear_logs")
112
+ )
113
+ return keyboard
114
+
115
+ def stream_command(chat_id, command):
116
+ initialize_logs_history(chat_id)
117
+ initialize_command_history(chat_id)
118
+ add_to_command_history(chat_id, command)
119
+
120
+ process = subprocess.Popen(
121
+ command,
122
+ stdout=subprocess.PIPE,
123
+ stderr=subprocess.STDOUT,
124
+ shell=True,
125
+ text=True,
126
+ bufsize=1,
127
+ universal_newlines=True
128
+ )
129
+
130
+ active_processes[chat_id] = {
131
+ 'process': process,
132
+ 'last_message_id': None,
133
+ 'buffer': "",
134
+ 'last_update': time.time(),
135
+ 'show_logs': True,
136
+ 'paused_at': None
137
+ }
138
+
139
+ UPDATE_INTERVAL = 2
140
+
141
+ try:
142
+ while True:
143
+ line = process.stdout.readline()
144
+ if not line and process.poll() is not None:
145
+ break
146
+
147
+ if line:
148
+ add_to_logs_history(chat_id, line.strip())
149
+ active_processes[chat_id]['buffer'] += line
150
+ current_time = time.time()
151
+
152
+ if (active_processes[chat_id]['show_logs'] and
153
+ current_time - active_processes[chat_id]['last_update'] >= UPDATE_INTERVAL):
154
+ send_log_update(chat_id)
155
+ active_processes[chat_id]['last_update'] = current_time
156
+
157
+ except Exception as e:
158
+ bot.send_message(chat_id, f"❌ Error: {str(e)}", parse_mode='HTML')
159
+ finally:
160
+ if chat_id in active_processes:
161
+ if active_processes[chat_id]['buffer'] and active_processes[chat_id]['show_logs']:
162
+ send_log_update(chat_id)
163
+
164
+ if process.poll() is not None:
165
+ bot.send_message(
166
+ chat_id,
167
+ f"βœ… Command completed with exit code: {process.returncode}",
168
+ parse_mode='HTML'
169
+ )
170
+ try:
171
+ del active_processes[chat_id]
172
+ except KeyError:
173
+ pass
174
+
175
+ def send_log_update(chat_id):
176
+ """Send or update log message"""
177
+ process_info = active_processes[chat_id]
178
+ log_content = process_info['buffer']
179
+
180
+ if process_info['process'].poll() is None:
181
+ chunks = split_logs(log_content)
182
+ total_pages = len(chunks)
183
+
184
+ # Determine the content of the current page (latest log chunk)
185
+ current_chunk = chunks[-1] if chunks else ""
186
+
187
+ # Add the page count to the log header
188
+ status = "πŸ“‹ Live Log (updating...)" if process_info['show_logs'] else "πŸ“‹ Live Log (paused)"
189
+ header = f"{status} (Page {total_pages}/{total_pages})\n"
190
+ formatted_message = f"{header}\n<pre>{sanitize_log(current_chunk)}</pre>"
191
+
192
+ if not process_info['show_logs']:
193
+ formatted_message += "\n<pre>Logs are paused. Click Resume to continue showing logs.</pre>"
194
+
195
+ # Append a timestamp to force a slight change in the content when paused
196
+ if not process_info['show_logs']:
197
+ formatted_message += f"\n\n<pre>Last paused at: {time.strftime('%H:%M:%S')}</pre>"
198
+
199
+ try:
200
+ if process_info['last_message_id']:
201
+ # Edit the current message only if content differs
202
+ bot.edit_message_text(
203
+ formatted_message,
204
+ chat_id=chat_id,
205
+ message_id=process_info['last_message_id'],
206
+ parse_mode='HTML',
207
+ reply_markup=create_log_keyboard()
208
+ )
209
+ else:
210
+ # Send a new message if there isn't an existing one
211
+ msg = bot.send_message(
212
+ chat_id,
213
+ formatted_message,
214
+ parse_mode='HTML',
215
+ reply_markup=create_log_keyboard()
216
+ )
217
+ process_info['last_message_id'] = msg.message_id
218
+ except telebot.apihelper.ApiException as e:
219
+ # Handle specific "message is not modified" errors
220
+ if "message is not modified" in str(e):
221
+ # Skip the update since content is the same
222
+ pass
223
+ else:
224
+ bot.send_message(chat_id, f"Error updating message: {str(e)}")
225
+ else:
226
+ chunks = split_logs(log_content)
227
+ for i, chunk in enumerate(chunks, 1):
228
+ formatted_chunk = (
229
+ f"πŸ“‹ Final Log Part {i}/{len(chunks)}\n"
230
+ f"<pre>{sanitize_log(chunk)}</pre>"
231
+ )
232
+ bot.send_message(chat_id, formatted_chunk, parse_mode='HTML')
233
+
234
+ if process_info['process'].poll() is not None:
235
+ process_info['buffer'] = ""
236
+
237
+
238
+ @bot.callback_query_handler(func=lambda call: True)
239
+ def handle_callback(call):
240
+ chat_id = call.message.chat.id
241
+ if chat_id not in active_processes and call.data != "clear_logs":
242
+ bot.answer_callback_query(call.id, "No active process")
243
+ return
244
+
245
+ if call.data == "stop_logs":
246
+ active_processes[chat_id]['show_logs'] = False
247
+ bot.answer_callback_query(call.id, "Logs paused")
248
+ send_log_update(chat_id)
249
+
250
+ elif call.data == "resume_logs":
251
+ active_processes[chat_id]['show_logs'] = True
252
+ bot.answer_callback_query(call.id, "Logs resumed")
253
+ send_log_update(chat_id)
254
+
255
+ elif call.data == "refresh_logs":
256
+ bot.answer_callback_query(call.id, "Logs refreshed")
257
+ send_log_update(chat_id)
258
+
259
+ elif call.data == "clear_logs":
260
+ if chat_id in logs_history:
261
+ logs_history[chat_id].clear()
262
+ if chat_id in active_processes:
263
+ active_processes[chat_id]['buffer'] = ""
264
+ bot.answer_callback_query(call.id, "Logs cleared")
265
+
266
+ @bot.message_handler(commands=['logs'])
267
+ def show_logs(message):
268
+ chat_id = message.chat.id
269
+
270
+ # Check if the user is authorized to use the bot
271
+ if not is_authorized(chat_id):
272
+ bot.reply_to(message, "❌ You are not authorized to use this bot.")
273
+ return
274
+
275
+ # Initialize the log history for the chat if it doesn't exist
276
+ initialize_logs_history(chat_id)
277
+
278
+ # Determine the logs content and whether it's live or historical
279
+ logs_content = ""
280
+ if chat_id in active_processes:
281
+ logs_content = active_processes[chat_id]['buffer']
282
+ is_live = True
283
+ else:
284
+ logs_content = get_latest_logs(chat_id)
285
+ is_live = False
286
+
287
+ # Show only the current live logs if the command is active
288
+ if is_live and logs_content:
289
+ # Get only the latest chunk of logs
290
+ chunks = split_logs(logs_content, latest_only=True)
291
+
292
+ if chunks:
293
+ status = "πŸ“‹ Live Logs"
294
+ formatted_chunk = f"{status} (Current Live Logs)\n<pre>{sanitize_log(chunks[0])}</pre>"
295
+
296
+ # Send the formatted log chunk as a message
297
+ msg = bot.send_message(
298
+ chat_id,
299
+ formatted_chunk,
300
+ parse_mode='HTML',
301
+ reply_markup=create_log_keyboard(is_live)
302
+ )
303
+
304
+ # Store the message ID for future updates
305
+ active_processes[chat_id]['last_message_id'] = msg.message_id
306
+ else:
307
+ # If there are logs to show and no active process
308
+ if logs_content:
309
+ # Split the logs into chunks to avoid hitting the message length limit
310
+ chunks = split_logs(logs_content)
311
+
312
+ # Send all parts of the logs
313
+ for i, chunk in enumerate(chunks, 1):
314
+ status = "πŸ“‹ Historical Logs"
315
+ formatted_chunk = f"{status} (Part {i}/{len(chunks)})\n<pre>{sanitize_log(chunk)}</pre>"
316
+
317
+ # Send the formatted log chunk as a message
318
+ bot.send_message(
319
+ chat_id,
320
+ formatted_chunk,
321
+ parse_mode='HTML',
322
+ reply_markup=create_log_keyboard(is_live)
323
+ )
324
+ else:
325
+ # If no logs are available, inform the user
326
+ bot.reply_to(message, "No logs available.")
327
+
328
+ @bot.message_handler(commands=['stop'])
329
+ def stop_command(message):
330
+ chat_id = message.chat.id
331
+
332
+ if not is_authorized(chat_id):
333
+ bot.reply_to(message, "❌ You are not authorized to use this bot.")
334
+ return
335
+
336
+ if chat_id in active_processes:
337
+ try:
338
+ # 1. First, disable log updates
339
+ active_processes[chat_id]['show_logs'] = False
340
+ active_processes[chat_id]['buffer'] = ""
341
+
342
+ # 2. Get the process
343
+ process = active_processes[chat_id]['process']
344
+
345
+ # 3. Try to terminate the process using Windows-friendly approach
346
+
347
+
348
+ try:
349
+ # Get the process using psutil
350
+ parent = psutil.Process(process.pid)
351
+ # Get all children processes
352
+ children = parent.children(recursive=True)
353
+
354
+ # Terminate children first
355
+ for child in children:
356
+ child.terminate()
357
+
358
+ # Terminate parent
359
+ parent.terminate()
360
+
361
+ # Wait for processes to terminate
362
+ psutil.wait_procs(children + [parent], timeout=3)
363
+
364
+ except:
365
+ # Fallback: try direct termination
366
+ try:
367
+ process.terminate()
368
+ except:
369
+ pass
370
+
371
+ # 4. Send success message
372
+ bot.reply_to(message, "πŸ›‘ Command stopped successfully.")
373
+
374
+ except Exception as e:
375
+ bot.reply_to(message, f"❌ Error stopping command: {str(e)}")
376
+
377
+ finally:
378
+ # 5. Clean up
379
+ if chat_id in active_processes:
380
+ try:
381
+ # One final attempt to kill if still running
382
+ process = active_processes[chat_id]['process']
383
+ if process.poll() is None:
384
+ process.kill()
385
+ except:
386
+ pass
387
+
388
+ # Remove from active processes
389
+ del active_processes[chat_id]
390
+ else:
391
+ bot.reply_to(message, "No active command to stop.")
392
+
393
+
394
+ @bot.message_handler(commands=['cmd'])
395
+ def show_command_history(message):
396
+ chat_id = message.chat.id
397
+
398
+ if not is_authorized(chat_id):
399
+ bot.reply_to(message, "❌ You are not authorized to use this bot.")
400
+ return
401
+
402
+ commands = get_command_history(chat_id)
403
+ if commands:
404
+ # formatted_commands = "\n".join(f"{i + 1}. {cmd}" for i, cmd in enumerate(commands))
405
+ formatted_commands = "\n".join([f"{i+1}. {cmd}" for i, cmd in enumerate(commands)])
406
+ # bot.reply_to(message, f"πŸ“œ Command History:\n{formatted_commands}")
407
+ # bot.reply_to(message, f"πŸ“œ Last {len(commands)} commands:\n<pre>{formatted_commands}</pre>", parse_mode='HTML')
408
+ bot.reply_to(message, f"πŸ“œ Command History (Last {len(commands)} commands):\n<pre>{formatted_commands}</pre>", parse_mode='HTML')
409
+
410
+ else:
411
+ bot.reply_to(message, "No command history available.")
412
+
413
+ @bot.message_handler(commands=['start'])
414
+ def send_welcome(message):
415
+ bot.reply_to(
416
+ message,
417
+ "πŸ‘‹ Welcome! Send me any command to execute and monitor logs.\n"
418
+ "Use /stop to stop the current command.\n"
419
+ "Use /cmds to see the last 20 executed commands."
420
+ )
421
+
422
+ @bot.message_handler(func=lambda message: True)
423
+ def execute_command(message):
424
+ chat_id = message.chat.id
425
+
426
+ if not is_authorized(chat_id):
427
+ bot.reply_to(message, "❌ You are not authorized to use this bot.")
428
+ return
429
+
430
+ if chat_id in active_processes:
431
+ bot.reply_to(
432
+ message,
433
+ "⚠️ A command is already running. Use /stop to stop it first."
434
+ )
435
+ return
436
+
437
+ command = message.text
438
+ bot.reply_to(message, f"▢️ Executing command: {command}")
439
+
440
+ thread = threading.Thread(
441
+ target=stream_command,
442
+ args=(chat_id, command)
443
+ )
444
+ thread.start()
445
+
446
+ # # Start the bot
447
+ # try:
448
+ # bot.polling(none_stop=True)
449
+ # except Exception as e:
450
+ # print(f"Bot polling error: {e}")
451
+
452
+
453
+ # ─────────────────────────────────────────────────── FastAPI ───
454
+ app = FastAPI()
455
+
456
+ @app.get("/")
457
+ def root(): # β‘‘ health‑check hits this β†’ must return 200 quickly
458
+ return {"status": "ok"}
459
+
460
+ @app.on_event("startup")
461
+ def startup():
462
+ # Launch the bot *after* Uvicorn has started
463
+ threading.Thread(target=bot.infinity_polling, daemon=True).start()