From c3cbcc1b2c9ad3d1dba570ad9cff2e6191c66e18 Mon Sep 17 00:00:00 2001
From: xaoyaoo <xaoyaoo@qq.com>
Date: Tue, 9 Jan 2024 22:02:11 +0800
Subject: [PATCH] =?UTF-8?q?=E6=9B=B4=E6=96=B0=E5=AF=BC=E5=87=BA=E4=B8=BAcs?=
 =?UTF-8?q?v=E7=9A=84=E6=96=B9=E5=BC=8F=EF=BC=8C=E5=AD=97=E6=AE=B5?=
 =?UTF-8?q?=E4=B8=AD=E6=9C=89=E5=9B=A0=E4=B8=BA,=E8=87=AA=E5=8A=A8?=
 =?UTF-8?q?=E8=BD=AC=E4=B9=89?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pywxdump/analyzer/export_chat.py |  7 +++--
 pywxdump/cli.py                  | 45 ++++++++++++++++----------------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/pywxdump/analyzer/export_chat.py b/pywxdump/analyzer/export_chat.py
index bffb001..92d0a32 100644
--- a/pywxdump/analyzer/export_chat.py
+++ b/pywxdump/analyzer/export_chat.py
@@ -12,6 +12,7 @@
 # Author:       xaoyaoo
 # Date:         2023/11/10
 # -------------------------------------------------------------------------------
+import csv
 import re
 import sqlite3
 import os
@@ -233,7 +234,9 @@ def export_csv(username, outpath, MSG_ALL_db_path, page_size=5000):
             break
         save_path = os.path.join(outpath, f"{username}_{i}_{i + page_size}.csv")
         with open(save_path, "w", encoding="utf-8") as f:
-            f.write("id,MsgSvrID,type_name,is_sender,talker,room_name,content,CreateTime\n")
+            csv_writer = csv.writer(f, quoting=csv.QUOTE_MINIMAL)
+            csv_writer.writerow(["id", "MsgSvrID", "type_name", "is_sender", "talker", "room_name", "content",
+                                 "CreateTime"])
             for row in data:
                 id = row.get("id", "")
                 MsgSvrID = row.get("MsgSvrID", "")
@@ -245,8 +248,8 @@ def export_csv(username, outpath, MSG_ALL_db_path, page_size=5000):
                 CreateTime = row.get("CreateTime", "")
 
                 content = json.dumps(content, ensure_ascii=False)
+                csv_writer.writerow([id, MsgSvrID, type_name, is_sender, talker, room_name, content, CreateTime])
 
-                f.write(f"{id},{MsgSvrID},{type_name},{is_sender},{talker},{room_name},{content},{CreateTime}\n")
     return True, f"导出成功: {outpath}"
 
 
diff --git a/pywxdump/cli.py b/pywxdump/cli.py
index 8ec2bf8..b999f8b 100644
--- a/pywxdump/cli.py
+++ b/pywxdump/cli.py
@@ -295,9 +295,7 @@ class MainExportChatRecords():
     def run(self, args):
         # 从命令行参数获取值
         t = args.type
-        if t not in ["html", "csv"]:
-            print("[-] 未知的导出类型")
-            return
+
         if t == "txt":
             try:
                 code, ret = export_csv(args.username, args.outpath, args.msg_path, page_size=10000000)
@@ -310,28 +308,31 @@ class MainExportChatRecords():
                 print(e)
                 print("[-] 导出失败")
                 return
+        elif t == "html":
+            try:
+                from flask import Flask, request, jsonify, render_template, g
+                import logging
+            except Exception as e:
+                print(e)
+                print("[-] 请安装flask( pip install flask)")
+                return
 
-        try:
-            from flask import Flask, request, jsonify, render_template, g
-            import logging
-        except Exception as e:
-            print(e)
-            print("[-] 请安装flask( pip install flask)")
+            if not os.path.exists(args.msg_path) or not os.path.exists(args.micro_path) or not os.path.exists(
+                    args.media_path):
+                print(os.path.exists(args.msg_path), os.path.exists(args.micro_path), os.path.exists(args.media_path))
+                print("[-] 输入数据库路径不存在")
+                return
+
+            if not os.path.exists(args.outpath):
+                os.makedirs(args.outpath)
+                print(f"[+] 创建输出文件夹：{args.outpath}")
+
+            export(args.username, args.outpath, args.msg_path, args.micro_path, args.media_path, args.filestorage_path)
+            print(f"[+] 导出成功{args.outpath}")
+        else:
+            print("[-] 未知的导出类型")
             return
 
-        if not os.path.exists(args.msg_path) or not os.path.exists(args.micro_path) or not os.path.exists(
-                args.media_path):
-            print(os.path.exists(args.msg_path), os.path.exists(args.micro_path), os.path.exists(args.media_path))
-            print("[-] 输入数据库路径不存在")
-            return
-
-        if not os.path.exists(args.outpath):
-            os.makedirs(args.outpath)
-            print(f"[+] 创建输出文件夹：{args.outpath}")
-
-        export(args.username, args.outpath, args.msg_path, args.micro_path, args.media_path, args.filestorage_path)
-        print(f"[+] 导出成功{args.outpath}")
-
 
 class MainAll():
     def init_parses(self, parser):