main.py 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200
  1. import asyncio
  2. from utils.config import get_config
  3. from utils.channel import (
  4. get_channel_items,
  5. append_data_to_info_data,
  6. append_total_data,
  7. sort_channel_list,
  8. write_channel_to_file,
  9. )
  10. from utils.tools import (
  11. update_file,
  12. get_pbar_remaining,
  13. get_ip_address,
  14. )
  15. from subscribe import get_channels_by_subscribe_urls
  16. from fofa import get_channels_by_fofa
  17. from online_search import get_channels_by_online_search
  18. import os
  19. from tqdm import tqdm
  20. from tqdm.asyncio import tqdm_asyncio
  21. from time import time
  22. from flask import Flask, render_template_string
  23. import sys
  24. config = get_config()
  25. app = Flask(__name__)
  26. @app.route("/")
  27. def show_result():
  28. user_final_file = getattr(config, "final_file", "result.txt")
  29. with open(user_final_file, "r", encoding="utf-8") as file:
  30. content = file.read()
  31. return render_template_string("<pre>{{ content }}</pre>", content=content)
  32. class UpdateSource:
  33. def __init__(self):
  34. self.run_ui = False
  35. self.tasks = []
  36. self.channel_items = get_channel_items()
  37. self.subscribe_result = {}
  38. self.multicast_result = {}
  39. self.online_search_result = {}
  40. self.channel_data = {}
  41. self.pbar = None
  42. self.total = 0
  43. self.start_time = None
  44. async def visit_page(self, channel_names=None):
  45. if config.open_subscribe:
  46. subscribe_task = asyncio.create_task(
  47. get_channels_by_subscribe_urls(self.update_progress)
  48. )
  49. self.tasks.append(subscribe_task)
  50. self.subscribe_result = await subscribe_task
  51. if config.open_multicast:
  52. multicast_task = asyncio.create_task(
  53. get_channels_by_fofa(self.update_progress)
  54. )
  55. self.tasks.append(multicast_task)
  56. self.multicast_result = await multicast_task
  57. if config.open_online_search:
  58. online_search_task = asyncio.create_task(
  59. get_channels_by_online_search(channel_names, self.update_progress)
  60. )
  61. self.tasks.append(online_search_task)
  62. self.online_search_result = await online_search_task
  63. def pbar_update(self, name=""):
  64. self.pbar.update()
  65. self.update_progress(
  66. f"正在进行{name}, 剩余{self.pbar.total - self.pbar.n}个接口, 预计剩余时间: {get_pbar_remaining(self.pbar, self.start_time)}",
  67. int((self.pbar.n / self.total) * 100),
  68. )
  69. async def main(self):
  70. try:
  71. channel_names = [
  72. name
  73. for channel_obj in self.channel_items.values()
  74. for name in channel_obj.keys()
  75. ]
  76. self.total = len(channel_names)
  77. await self.visit_page(channel_names)
  78. self.tasks = []
  79. self.channel_data = append_total_data(
  80. self.channel_items.items(),
  81. self.channel_data,
  82. self.subscribe_result,
  83. self.multicast_result,
  84. self.online_search_result,
  85. )
  86. if config.open_sort:
  87. semaphore = asyncio.Semaphore(100)
  88. self.tasks = [
  89. asyncio.create_task(
  90. sort_channel_list(
  91. semaphore,
  92. cate,
  93. name,
  94. info_list,
  95. lambda: self.pbar_update("测速排序"),
  96. )
  97. )
  98. for cate, channel_obj in self.channel_data.items()
  99. for name, info_list in channel_obj.items()
  100. ]
  101. self.pbar = tqdm_asyncio(total=len(self.tasks), desc="Sorting")
  102. self.update_progress(
  103. f"正在测速排序, 共{len(self.tasks)}个频道",
  104. int((self.pbar.n / len(self.tasks)) * 100),
  105. )
  106. self.start_time = time()
  107. sort_results = await tqdm_asyncio.gather(*self.tasks, disable=True)
  108. self.channel_data = {}
  109. for result in sort_results:
  110. if result:
  111. cate = result.get("cate")
  112. name = result.get("name")
  113. data = result.get("data")
  114. self.channel_data = append_data_to_info_data(
  115. self.channel_data, cate, name, data, False
  116. )
  117. self.pbar = tqdm(total=self.total, desc="Writing")
  118. self.start_time = time()
  119. write_channel_to_file(
  120. self.channel_items.items(),
  121. self.channel_data,
  122. lambda: self.pbar_update("写入结果"),
  123. )
  124. self.pbar.close()
  125. user_final_file = getattr(config, "final_file", "result.txt")
  126. update_file(user_final_file, "result_new.txt")
  127. if config.open_sort:
  128. user_log_file = (
  129. "user_result.log"
  130. if os.path.exists("user_config.py")
  131. else "result.log"
  132. )
  133. update_file(user_log_file, "result_new.log")
  134. print(f"Update completed! Please check the {user_final_file} file!")
  135. if not os.environ.get("GITHUB_ACTIONS"):
  136. print(f"You can access the result at {get_ip_address()}")
  137. if self.run_ui:
  138. self.update_progress(
  139. f"更新完成, 请检查{user_final_file}文件, 可访问以下链接:",
  140. 100,
  141. True,
  142. url=f"{get_ip_address()}",
  143. )
  144. except asyncio.exceptions.CancelledError:
  145. print("Update cancelled!")
  146. async def start(self, callback=None):
  147. def default_callback(self, *args, **kwargs):
  148. pass
  149. self.update_progress = callback or default_callback
  150. self.run_ui = True if callback else False
  151. if config.open_update:
  152. await self.main()
  153. if self.run_ui:
  154. if not config.open_update:
  155. print(f"You can access the result at {get_ip_address()}")
  156. self.update_progress(
  157. f"服务启动成功, 可访问以下链接:",
  158. 100,
  159. True,
  160. url=f"{get_ip_address()}",
  161. )
  162. app.run(host="0.0.0.0", port=8000)
  163. def stop(self):
  164. for task in self.tasks:
  165. task.cancel()
  166. self.tasks = []
  167. if self.pbar:
  168. self.pbar.close()
  169. def scheduled_task():
  170. if config.open_update:
  171. update_source = UpdateSource()
  172. loop = asyncio.new_event_loop()
  173. asyncio.set_event_loop(loop)
  174. loop.run_until_complete(update_source.start())
  175. if __name__ == "__main__":
  176. # Run scheduled_task
  177. scheduled_task()
  178. # If not run with 'scheduled_task' argument and not in GitHub Actions, start Flask server
  179. if len(sys.argv) <= 1 or sys.argv[1] != "scheduled_task":
  180. if not os.environ.get("GITHUB_ACTIONS"):
  181. app.run(host="0.0.0.0", port=3000)