From 96d14f39f60a239c2e9ceee1f742502fdecc4aa8 Mon Sep 17 00:00:00 2001
From: ivy1996-encode <15055461510@163.com>
Date: Sun, 14 Mar 2021 12:12:33 +0800
Subject: [PATCH] add message file download support

---
 bot.go      |  2 +-
 bot_test.go | 15 ++++++++++++---
 client.go   | 31 +++++++++++++++++++++++++++++++
 global.go   |  5 ++++-
 message.go  | 26 +++++++++++++++++++++-----
 5 files changed, 69 insertions(+), 10 deletions(-)

diff --git a/bot.go b/bot.go
index 7f96399..fcc64c8 100644
--- a/bot.go
+++ b/bot.go
@@ -193,7 +193,7 @@ func (b *Bot) Block() error {
 	if b.self == nil {
 		return errors.New("`Block` must be called after user login")
 	}
-	if _, closed := <-b.exit; closed {
+	if _, closed := <-b.exit; !closed {
 		return errors.New("can not call `Block` after user logout")
 	}
 	close(b.exit)
diff --git a/bot_test.go b/bot_test.go
index 2eb6ce9..404a023 100644
--- a/bot_test.go
+++ b/bot_test.go
@@ -2,22 +2,31 @@ package openwechat
 
 import (
 	"fmt"
+	"io/ioutil"
 	"testing"
 )
 
 func TestDefaultBot(t *testing.T) {
 	bot := DefaultBot()
 	messageHandler := func(message *Message) {
-		fmt.Println(message.Content)
+		if message.HasFile() {
+			if message.IsMedia() {
+				resp, err := message.GetFile()
+				if err == nil {
+					data, _ := ioutil.ReadAll(resp.Body)
+					ioutil.WriteFile(message.EncryFileName, data, 0x777)
+					resp.Body.Close()
+				}
+			}
+		}
 	}
 	bot.RegisterMessageHandler(messageHandler)
 	bot.UUIDCallback = PrintlnQrcodeUrl
 	if err := bot.Login(); err != nil {
 		fmt.Println(err)
 		return
-
 	}
-	bot.Block()
+	fmt.Println(bot.Block())
 }
 
 func TestBotMessageHandler(t *testing.T) {
diff --git a/client.go b/client.go
index a5ae9df..f0708f8 100644
--- a/client.go
+++ b/client.go
@@ -359,3 +359,34 @@ func (c *Client) WebWxGetMsgImg(msg *Message, info LoginInfo) (*http.Response, e
 	path.RawQuery = params.Encode()
 	return c.Get(path.String())
 }
+
+func (c *Client) WebWxGetVoice(msg *Message, info LoginInfo) (*http.Response, error) {
+	path, _ := url.Parse(webWxGetVoiceUrl)
+	params := url.Values{}
+	params.Add("msgid", msg.MsgId)
+	params.Add("skey", info.SKey)
+	path.RawQuery = params.Encode()
+	return c.Get(path.String())
+}
+
+func (c *Client) WebWxGetVideo(msg *Message, info LoginInfo) (*http.Response, error) {
+	path, _ := url.Parse(webWxGetVideoUrl)
+	params := url.Values{}
+	params.Add("msgid", msg.MsgId)
+	params.Add("skey", info.SKey)
+	path.RawQuery = params.Encode()
+	return c.Get(path.String())
+}
+
+func (c *Client) WebWxGetMedia(msg *Message, info LoginInfo) (*http.Response, error) {
+	path, _ := url.Parse(webWxGetMediaUrl)
+	params := url.Values{}
+	params.Add("sender", msg.FromUserName)
+	params.Add("mediaid", msg.MediaId)
+	params.Add("encryfilename", msg.EncryFileName)
+	params.Add("fromuser", fmt.Sprintf("%d", info.WxUin))
+	params.Add("pass_ticket", info.PassTicket)
+	params.Add("webwx_data_ticket", getWebWxDataTicket(c.Jar.Cookies(path)))
+	path.RawQuery = params.Encode()
+	return c.Get(path.String())
+}
diff --git a/global.go b/global.go
index 1381183..15487a2 100644
--- a/global.go
+++ b/global.go
@@ -29,7 +29,10 @@ const (
 	webWxVerifyUserUrl      = "https://wx2.qq.com/cgi-bin/mmwebwx-bin/webwxverifyuser"
 	syncCheckUrl            = "https://webpush.wx2.qq.com/cgi-bin/mmwebwx-bin/synccheck"
 	webWxUpLoadMediaUrl     = "https://file.wx2.qq.com/cgi-bin/mmwebwx-bin/webwxuploadmedia"
-	webWxGetMsgImgUrl          = "https://wx2.qq.com/cgi-bin/mmwebwx-bin/webwxgetmsgimg"
+	webWxGetMsgImgUrl       = "https://wx2.qq.com/cgi-bin/mmwebwx-bin/webwxgetmsgimg"
+	webWxGetVoiceUrl        = "https://wx2.qq.com/cgi-bin/mmwebwx-bin/webwxgetvoice"
+	webWxGetVideoUrl        = "https://wx2.qq.com/cgi-bin/mmwebwx-bin/webwxgetvideo"
+	webWxGetMediaUrl        = "https://file.wx2.qq.com/cgi-bin/mmwebwx-bin/webwxgetmedia"
 
 	jsonContentType = "application/json; charset=utf-8"
 )
diff --git a/message.go b/message.go
index ee8bba1..b219fcc 100644
--- a/message.go
+++ b/message.go
@@ -159,7 +159,7 @@ func (m *Message) IsVideo() bool {
 	return m.MsgType == 43 || m.MsgType == 62
 }
 
-func (m *Message) IsSharing() bool {
+func (m *Message) IsMedia() bool {
 	return m.MsgType == 49
 }
 
@@ -175,11 +175,27 @@ func (m *Message) IsNotify() bool {
 	return m.MsgType == 51 && m.StatusNotifyCode != 0
 }
 
-func (m *Message) GetMsgImageResponse() (*http.Response, error) {
-	if !m.IsPicture() {
-		return nil, errors.New("Picture type message required")
+func (m *Message) HasFile() bool {
+	return m.IsPicture() || m.IsVoice() || m.IsVideo() || m.IsMedia()
+}
+
+func (m *Message) GetFile() (*http.Response, error) {
+	if !m.HasFile() {
+		return nil, errors.New("invalid message type")
 	}
-	return m.Bot.Caller.Client.WebWxGetMsgImg(m, m.Bot.storage.GetLoginInfo())
+	if m.IsPicture() {
+		return m.Bot.Caller.Client.WebWxGetMsgImg(m, m.Bot.storage.GetLoginInfo())
+	}
+	if m.IsVoice() {
+		return m.Bot.Caller.Client.WebWxGetVoice(m, m.Bot.storage.GetLoginInfo())
+	}
+	if m.IsVideo() {
+		return m.Bot.Caller.Client.WebWxGetVideo(m, m.Bot.storage.GetLoginInfo())
+	}
+	if m.IsMedia() {
+		return m.Bot.Caller.Client.WebWxGetMedia(m, m.Bot.storage.GetLoginInfo())
+	}
+	return nil, errors.New("unsupported type")
 }
 
 // 用在多个messageHandler之间传递信息