intro pymongo

messages/base.py

@@ -1,13 +1,45 @@
+from store.mongo import MongoStore
+
 class Base:
 
     instance = None
 
     def set_payload(self, payload):
         self.instance.ParseFromString(payload)
-    
+
     def user(self):
         return self.instance.user
 
+    def persists(self):
+        try:
+            store = MongoStore()
+            store.set_collection('user')
+
+            user = self.user()
+
+            store.insert_one({
+                "id": user.id,
+                "shortId": user.shortId,
+                "nickname": user.nickname,
+                "gender": user.gender,
+                "avatar_thumb": user.avatarThumb.urlList[0],
+                "followInfo": {
+                    "followingCount": user.followInfo.followingCount,
+                    "followerCount": user.followInfo.followerCount
+                }
+            })
+
+            store.set_collection(self.instance.common.method)
+            store.insert_one({
+                "msgId": self.instance.common.msgId,
+                "roomId": self.instance.common.roomId,
+                "userId": user.id,
+                'content': self.format_content()
+            })
+        except Exception as e:
+            print(self.instance.common.method + 'persists error')
+
+
     def __str__(self):
         pass
 

messages/chat.py

@@ -7,5 +7,8 @@ class ChatMessage(Base):
     def __init__(self):
         self.instance = message_pb2.ChatMessage()
 
+    def format_content(self):
+        return self.user().nickname + ': ' + self.instance.content
+
     def __str__(self):
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【发言】' + self.user().nickname + ': ' + self.instance.content
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【发言】' +  self.format_content()

messages/gift.py

@@ -7,5 +7,8 @@ class GiftMessage(Base):
     def __init__(self):
         self.instance = message_pb2.GiftMessage()
 
+    def format_content(self):
+        return self.instance.common.describe
+
     def __str__(self):
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【送礼】' + self.instance.common.describe
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【送礼】' + self.format_content()

messages/like.py

@@ -7,5 +7,8 @@ class LikeMessage(Base):
     def __init__(self):
         self.instance = message_pb2.LikeMessage()
 
+    def format_content(self):
+        return self.user().nickname + ' 点赞了直播间（' + str(self.instance.count) + '连赞）'
+
     def __str__(self):
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【点赞】' +  self.user().nickname + ' 点赞了直播间（' + str(self.instance.count) + '连赞）'
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【点赞】' +  self.format_content()

messages/member.py

@@ -7,8 +7,12 @@ class MemberMessage(Base):
     def __init__(self):
         self.instance = message_pb2.MemberMessage()
 
-    def __str__(self):
+
+    def format_content(self):
         template = self.instance.common.displayText.defaultPattern
         nickname = self.user().nickname
 
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【进入直播间】' + template.replace('{0:user}', nickname).replace('{1:string}', '')
+        return template.replace('{0:user}', nickname).replace('{1:string}', '')
+
+    def __str__(self):
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【进入直播间】' + self.format_content()

messages/roomuserseq.py

@@ -7,5 +7,8 @@ class RoomUserSeqMessage(Base):
     def __init__(self):
         self.instance = message_pb2.RoomUserSeqMessage()
 
+    def format_content(self):
+        return self.instance.totalUserStr
+
     def __str__(self):
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【观看人数】' + self.instance.totalUserStr
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【观看人数】' + self.format_content()

messages/social.py

@@ -7,5 +7,8 @@ class SocialMessage(Base):
     def __init__(self):
         self.instance = message_pb2.SocialMessage()
 
+    def format_content(self):
+        return self.user().nickname + ' 关注了主播'
+
     def __str__(self):
-        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【关注】' +  self.user().nickname + ' 关注了主播'
+        return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) + '【关注】' +  self.format_content()

messages/utils.py

@@ -36,31 +36,37 @@ def decodeMsg(messages):
             if message.method == 'WebcastMemberMessage':
                 member_message = MemberMessage()
                 member_message.set_payload(message.payload)
+                member_message.persists()
                 print(f"\n{RED}[+] {member_message} {RESET}")
 
             elif message.method == 'WebcastSocialMessage':
                 social_message = SocialMessage()
                 social_message.set_payload(message.payload)
+                social_message.persists()
                 print(f"\n{GREEN}[+] {social_message} {RESET}")
 
             elif message.method == 'WebcastChatMessage':
                 chat_message = ChatMessage()
                 chat_message.set_payload(message.payload)
+                chat_message.persists()
                 print(f"\n{BLUE}[+] {chat_message} {RESET}")
 
             elif message.method == 'WebcastLikeMessage':
                 like_message = LikeMessage()
                 like_message.set_payload(message.payload)
+                like_message.persists()
                 print(f"\n{CYAN}[+] {like_message} {RESET}")
 
             elif message.method == 'WebcastGiftMessage':
                 gift_message = GiftMessage()
                 gift_message.set_payload(message.payload)
+                gift_message.persists()
                 print(f"\n{MAGENTA}[+] {gift_message} {RESET}")
 
             elif message.method == 'WebcastRoomUserSeqMessage':
                 room_user_seq_message = RoomUserSeqMessage() 
                 room_user_seq_message.set_payload(message.payload)
+                # room_user_seq_message.persists()
                 print(f"\n{YELLOW}[+] {room_user_seq_message} {RESET}")
 
         except Exception as e:

protobuf/message.proto

@@ -419,6 +419,22 @@ message TextPiece{
 }
 
 message Image{
+    repeated string urlList = 1;
+    string uri = 2;
+    int64 height = 3;
+    int64 width = 4;
+    string avgColor = 5;
+    int32 imageType = 6;
+    string openWebUrl = 7;
+    Content content = 8;
+    bool isAnimated = 9;
+
+    message Content {
+        string name = 1;
+        string fontColor = 2;
+        int64 level = 3;
+        string alternativeText = 4;
+    }
 }
 
 message TextPieceUser{

requirements.txt

@@ -1,3 +1,4 @@
 colorama==0.4.4
 protobuf==3.19.1
+pymongo==3.12.1
 watchdog==2.1.6

store/mongo.py

@@ -0,0 +1,15 @@
+import pymongo
+
+class MongoStore:
+    def __init__(self):
+        self.client = pymongo.MongoClient("mongodb://localhost:27017/")
+        self.db = self.client['tiktok']
+    
+    def set_collection(self, collection):
+        self.collection = self.db[collection]
+    
+    def insert_one(self, data):
+        return self.collection.insert_one(data)
+
+    def insert_many(self, data):
+        return self.collection.insert_many(data)