intro pyyml

config/helper.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+import yaml
+
+def config():
+    settings_file = str(Path(__file__).parent.absolute()) + '/settings.yml'
+
+    with open(settings_file, 'r') as f:
+        return yaml.load(f, Loader=yaml.FullLoader)
+

config/settings.yml

@@ -0,0 +1,11 @@
+watchdog:
+  dir: '/Users/geng/douyin_live/'
+
+webdriver:
+  bin: '/usr/local/bin/chromedriver'
+  proxy: '127.0.0.1:8080'
+  
+mongo:
+  uri : 'mongodb://localhost:27017/'
+  dbname: 'tiktok'
+  enabled: 0

main.py

@@ -1,9 +1,15 @@
 import sys
+from urllib.parse import urlparse
 
 from scripts import watcher, webdriver
+from config.helper import config
 
 if __name__ == '__main__':
-    webdriver.lunch(sys.argv[1])
+    if len(sys.argv) == 1 or not urlparse(sys.argv[1]).scheme:
+        print('Invalid url provided, please check...')
+        sys.exit(1)
 
-    w = watcher.Watcher(directory='/Users/geng/douyin_live')
+    webdriver.go(sys.argv[1])
+
+    w = watcher.Watcher(directory=config()['watchdog']['dir'])
     w.run()

messages/base.py

@@ -1,5 +1,7 @@
 from store.mongo import MongoStore
 
+from config.helper import config
+
 class Base:
 
     instance = None
@@ -11,6 +13,9 @@ class Base:
         return self.instance.user
 
     def persists(self):
+        if config()['mongo']['enabled'] != 'on':
+            return
+
         try:
             store = MongoStore()
             store.set_collection('user')

requirements.txt

@@ -11,6 +11,7 @@ protobuf==3.19.1
 pycparser==2.21
 pymongo==3.12.1
 pyOpenSSL==21.0.0
+PyYAML==6.0
 selenium==4.1.0
 six==1.16.0
 sniffio==1.2.0

scripts/mitmproxy.py

@@ -4,10 +4,12 @@
 import uuid
 from mitmproxy import http
 
+from config.helper import config
+
 class Writer:
     def response(self, flow: http.HTTPFlow) -> None:
         if flow.request.host == 'live.douyin.com':
-            with open('/Users/geng/douyin_live/' + uuid.uuid4().hex, 'wb') as f:
+            with open(config().mitmproxy.log_dir + uuid.uuid4().hex, 'wb') as f:
                 f.write(bytes(flow.response.content))
 
 addons = [Writer()]

scripts/webdriver.py

@@ -1,13 +1,24 @@
 from selenium import webdriver
 from selenium.webdriver.chrome.options import Options
+from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
+from selenium.webdriver.common.proxy import Proxy, ProxyType
 
-def lunch(url):
+from config.helper import config
+
+def go(url):
     chrome_options = Options()
-    chrome_options.add_argument('--proxy-server=127.0.0.1:8080')
+    chrome_options.add_argument('--proxy-server=%s' % config()['webdriver']['proxy'])
-    chrome_options.add_argument('--headless')
+    # chrome_options.add_argument('--headless')
     chrome_options.add_experimental_option('detach', True)
 
-    driver = webdriver.Chrome(options=chrome_options)
+    proxy = Proxy()
+    proxy.proxy_type = ProxyType.MANUAL
+    proxy.http_proxy = config()['webdriver']['proxy']
+    proxy.ssl_proxy = config()['webdriver']['proxy']
+
+    capabilities = DesiredCapabilities.CHROME
+    proxy.add_to_capabilities(capabilities)
+
+    driver = webdriver.Chrome(options=chrome_options, desired_capabilities=capabilities, executable_path=config()['webdriver']['bin'])
 
     driver.get(url)
-

store/mongo.py

@@ -1,9 +1,11 @@
 import pymongo
 
+from config.helper import config
+
 class MongoStore:
     def __init__(self):
-        self.client = pymongo.MongoClient("mongodb://localhost:27017/")
+        self.client = pymongo.MongoClient(config()['mongo']['uri'])
-        self.db = self.client['tiktok']
+        self.db = self.client[config()['mongo']['dbname']]
     
     def set_collection(self, collection):
         self.collection = self.db[collection]