understanding commited on
Commit
fc8fccc
·
verified ·
1 Parent(s): 2df54ff

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +40 -78
main.py CHANGED
@@ -14,79 +14,43 @@ logger = logging.getLogger(__name__)
14
  # Retrieve Session String from Hugging Face secrets
15
  SESSION_STRING = os.environ.get('SESSION_STRING')
16
 
17
- # --- Dynamic Storage Path Initialization ---
18
- # Define potential base directories for storage and a subdirectory name
19
- PERSISTENT_STORAGE_CANDIDATES = [
20
- Path("/data"), # Common absolute path for persistent data in containers
21
- Path("."), # Current working directory (e.g., /app if WORKDIR is /app)
22
- ]
23
- APP_STORAGE_SUBDIR_NAME = "app_template_storage" # Subdirectory to hold our template
24
  TEMPLATE_FILENAME = "user_template_content.txt"
25
 
26
- DATA_DIR: Path = None
27
- TEMPLATE_FILE_PATH: Path = None
 
 
28
 
29
- def initialize_storage_paths():
30
- """Tries to find/create a writable directory for the template file."""
31
- global DATA_DIR, TEMPLATE_FILE_PATH
32
 
33
- # 1. Try creating a subdirectory in candidate locations
34
- for base_candidate in PERSISTENT_STORAGE_CANDIDATES:
35
- potential_data_dir = base_candidate / APP_STORAGE_SUBDIR_NAME
36
- try:
37
- potential_data_dir.mkdir(parents=True, exist_ok=True)
38
- # Test if we can actually write a dummy file here
39
- test_file = potential_data_dir / ".writable_test"
40
- with open(test_file, "w") as f:
41
- f.write("test")
42
- test_file.unlink() # Clean up test file
43
-
44
- DATA_DIR = potential_data_dir
45
- TEMPLATE_FILE_PATH = DATA_DIR / TEMPLATE_FILENAME
46
- logger.info(f"Successfully established writable data directory: {DATA_DIR}")
47
- return
48
- except PermissionError:
49
- logger.warning(f"Permission denied to create/use directory {potential_data_dir}. Trying next candidate.")
50
- except Exception as e:
51
- logger.warning(f"Failed to create/use directory {potential_data_dir} due to {type(e).__name__}: {e}. Trying next candidate.")
52
-
53
- # 2. If subdirectory creation failed, try to place the file directly in a candidate base directory
54
- logger.warning("Could not create/use a dedicated subdirectory. Attempting to use a base directory directly for the file.")
55
- for base_candidate in PERSISTENT_STORAGE_CANDIDATES:
56
- potential_template_file_path = base_candidate / TEMPLATE_FILENAME
57
- try:
58
- # Test writability by trying to open the file in append mode (creates if not exists)
59
- with open(potential_template_file_path, 'a'):
60
- pass # Just testing if open works without error
61
- # If we are here, it means we can probably write the file directly.
62
- DATA_DIR = base_candidate # The "data directory" is now the base itself
63
- TEMPLATE_FILE_PATH = potential_template_file_path
64
- logger.info(f"Using base directory {DATA_DIR.resolve()} directly for template file: {TEMPLATE_FILE_PATH}")
65
- return
66
- except PermissionError:
67
- logger.warning(f"Permission denied to write template file directly in {base_candidate.resolve()}. Trying next candidate.")
68
- except Exception as e:
69
- logger.warning(f"Failed to write template file in {base_candidate.resolve()} due to {type(e).__name__}: {e}. Trying next candidate.")
70
-
71
- # If all attempts fail
72
- logger.error("CRITICAL: Unable to find or create a writable location for the template file.")
73
- # Set paths to None to indicate failure; the app might crash or malfunction later.
74
- # A more robust solution would be to exit here.
75
- DATA_DIR = None
76
- TEMPLATE_FILE_PATH = None
77
- # Consider exiting: exit("Failed to initialize storage.")
78
-
79
- # Call this ONCE at the start of the script
80
- initialize_storage_paths()
81
-
82
- # Global variable to hold the loaded template content
83
- current_template_content = None
84
 
85
  # --- Helper Functions for Template Management ---
86
  async def load_template_from_file():
87
  """Loads template content from the persistent file."""
88
  global current_template_content
89
- if not TEMPLATE_FILE_PATH:
90
  logger.error("Template file path is not configured. Cannot load template.")
91
  current_template_content = None
92
  return
@@ -106,15 +70,12 @@ async def load_template_from_file():
106
  async def save_template_to_file(content: str):
107
  """Saves template content to the persistent file."""
108
  global current_template_content
109
- if not TEMPLATE_FILE_PATH:
110
  logger.error("Template file path is not configured. Cannot save template.")
111
  return False
112
-
113
  try:
114
- # Ensure the parent directory exists (it should if initialize_storage_paths worked for a subdir)
115
- if TEMPLATE_FILE_PATH.parent != Path("."): # Avoid trying to create "." if file is in current dir
116
- TEMPLATE_FILE_PATH.parent.mkdir(parents=True, exist_ok=True)
117
-
118
  with open(TEMPLATE_FILE_PATH, "w", encoding="utf-8") as f:
119
  f.write(content)
120
  current_template_content = content
@@ -129,15 +90,15 @@ if not SESSION_STRING:
129
  logger.error("SESSION_STRING environment variable not found. Please set it in Hugging Face secrets.")
130
  exit(1)
131
 
132
- if not TEMPLATE_FILE_PATH:
133
- logger.critical("CRITICAL: Template storage path could not be initialized. The application cannot manage templates. Exiting.")
 
134
  exit(1)
135
 
136
-
137
  logger.info("Initializing Hydrogram Client with session string...")
138
  try:
139
  app = Client(
140
- name="user_session_hgs", # Changed name slightly to avoid old session file conflicts if any
141
  session_string=SESSION_STRING,
142
  )
143
  logger.info("Hydrogram Client initialized.")
@@ -146,13 +107,14 @@ except Exception as e:
146
  exit(1)
147
 
148
 
149
- # --- Bot Event Handlers (Unchanged from your previous version) ---
150
  @app.on_message(filters.command("start") & filters.private)
151
  async def start_handler(client: Client, message: Message):
152
  sender_name = message.from_user.first_name if message.from_user else "User"
153
  logger.info(f"Received /start command from {sender_name} (ID: {message.from_user.id})")
154
  welcome_text = f"Hello {sender_name}!\n"
155
  welcome_text += "I am ready to manage your template.\n"
 
156
  welcome_text += "Use /settemplate <your text> to set a new template.\n"
157
  welcome_text += "Use /gettemplate to view the current template."
158
  if current_template_content:
@@ -187,7 +149,7 @@ async def get_template_handler(client: Client, message: Message):
187
  response_text = "ℹ️ No template is currently set. Use `/settemplate <your text>` to set one."
188
  await message.reply_text(response_text)
189
 
190
- # --- Main Execution ---
191
  async def main():
192
  await load_template_from_file()
193
  logger.info("Attempting to connect and start the bot...")
@@ -204,7 +166,7 @@ async def main():
204
  except Exception as e:
205
  logger.error(f"An unexpected error occurred during bot startup or runtime: {type(e).__name__} - {e}", exc_info=True)
206
  finally:
207
- if app.is_initialized and app.is_connected: # Check if initialized before checking is_connected
208
  logger.info("Stopping the bot...")
209
  await app.stop()
210
  logger.info("Bot stopped.")
@@ -216,7 +178,7 @@ if __name__ == '__main__':
216
  asyncio.run(main())
217
  except KeyboardInterrupt:
218
  logger.info("Bot manually interrupted. Exiting...")
219
- except SystemExit as e: # Catch explicit exits
220
  logger.info(f"Application exiting with status: {e.code}")
221
  except Exception as e:
222
  logger.critical(f"Critical error in main execution block: {type(e).__name__} - {e}", exc_info=True)
 
14
  # Retrieve Session String from Hugging Face secrets
15
  SESSION_STRING = os.environ.get('SESSION_STRING')
16
 
17
+ # --- Storage Path Initialization (Aligned with Dockerfile) ---
18
+ # This name MUST match the directory created in the Dockerfile
19
+ # (either the default in ARG or how you set it in `RUN mkdir -p ./${APP_TEMPLATE_DIR_NAME}`)
20
+ APP_TEMPLATE_SUBDIR_NAME = "app_template_storage"
 
 
 
21
  TEMPLATE_FILENAME = "user_template_content.txt"
22
 
23
+ # The data directory is now relative to the WORKDIR (e.g., /app/app_template_storage)
24
+ # This path is prepared by the Dockerfile.
25
+ DATA_DIR = Path(".") / APP_TEMPLATE_SUBDIR_NAME
26
+ TEMPLATE_FILE_PATH = DATA_DIR / TEMPLATE_FILENAME
27
 
28
+ current_template_content = None # Global variable for template
 
 
29
 
30
+ def verify_storage_path_is_writable():
31
+ """Verifies that the pre-configured DATA_DIR is writable."""
32
+ try:
33
+ # The Dockerfile should have created DATA_DIR.
34
+ # Python's mkdir with exist_ok=True is safe even if it exists.
35
+ # More importantly, we test if we can write into it.
36
+ DATA_DIR.mkdir(parents=True, exist_ok=True)
37
+
38
+ test_file = DATA_DIR / ".writable_test_py"
39
+ with open(test_file, "w") as f:
40
+ f.write("test")
41
+ test_file.unlink() # Clean up test file
42
+ logger.info(f"Successfully verified writable data directory: {DATA_DIR.resolve()}")
43
+ return True
44
+ except Exception as e:
45
+ logger.error(f"CRITICAL: Data directory {DATA_DIR.resolve()} (expected to be created and permissioned by Dockerfile) is not writable or accessible.")
46
+ logger.error(f"Error details: {type(e).__name__} - {e}")
47
+ return False
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
 
49
  # --- Helper Functions for Template Management ---
50
  async def load_template_from_file():
51
  """Loads template content from the persistent file."""
52
  global current_template_content
53
+ if not TEMPLATE_FILE_PATH: # Should not happen if verify_storage_path_is_writable passes
54
  logger.error("Template file path is not configured. Cannot load template.")
55
  current_template_content = None
56
  return
 
70
  async def save_template_to_file(content: str):
71
  """Saves template content to the persistent file."""
72
  global current_template_content
73
+ if not TEMPLATE_FILE_PATH: # Should not happen
74
  logger.error("Template file path is not configured. Cannot save template.")
75
  return False
 
76
  try:
77
+ # DATA_DIR.mkdir(parents=True, exist_ok=True) # Dockerfile should handle this.
78
+ # verify_storage_path_is_writable also calls it.
 
 
79
  with open(TEMPLATE_FILE_PATH, "w", encoding="utf-8") as f:
80
  f.write(content)
81
  current_template_content = content
 
90
  logger.error("SESSION_STRING environment variable not found. Please set it in Hugging Face secrets.")
91
  exit(1)
92
 
93
+ # Verify the storage path *before* initializing the client or doing anything else
94
+ if not verify_storage_path_is_writable():
95
+ logger.critical("Exiting due to storage path issues. Check Dockerfile and permissions.")
96
  exit(1)
97
 
 
98
  logger.info("Initializing Hydrogram Client with session string...")
99
  try:
100
  app = Client(
101
+ name="user_session_hgs_v3", # Slightly new name again
102
  session_string=SESSION_STRING,
103
  )
104
  logger.info("Hydrogram Client initialized.")
 
107
  exit(1)
108
 
109
 
110
+ # --- Bot Event Handlers (Unchanged) ---
111
  @app.on_message(filters.command("start") & filters.private)
112
  async def start_handler(client: Client, message: Message):
113
  sender_name = message.from_user.first_name if message.from_user else "User"
114
  logger.info(f"Received /start command from {sender_name} (ID: {message.from_user.id})")
115
  welcome_text = f"Hello {sender_name}!\n"
116
  welcome_text += "I am ready to manage your template.\n"
117
+ welcome_text += "Template storage is at: " + str(TEMPLATE_FILE_PATH.resolve()) + "\n"
118
  welcome_text += "Use /settemplate <your text> to set a new template.\n"
119
  welcome_text += "Use /gettemplate to view the current template."
120
  if current_template_content:
 
149
  response_text = "ℹ️ No template is currently set. Use `/settemplate <your text>` to set one."
150
  await message.reply_text(response_text)
151
 
152
+ # --- Main Execution (Unchanged) ---
153
  async def main():
154
  await load_template_from_file()
155
  logger.info("Attempting to connect and start the bot...")
 
166
  except Exception as e:
167
  logger.error(f"An unexpected error occurred during bot startup or runtime: {type(e).__name__} - {e}", exc_info=True)
168
  finally:
169
+ if app.is_initialized and app.is_connected:
170
  logger.info("Stopping the bot...")
171
  await app.stop()
172
  logger.info("Bot stopped.")
 
178
  asyncio.run(main())
179
  except KeyboardInterrupt:
180
  logger.info("Bot manually interrupted. Exiting...")
181
+ except SystemExit as e:
182
  logger.info(f"Application exiting with status: {e.code}")
183
  except Exception as e:
184
  logger.critical(f"Critical error in main execution block: {type(e).__name__} - {e}", exc_info=True)