From 0487e16623bdcb03c5041a91699c481b38687b3a Mon Sep 17 00:00:00 2001 From: Luke Harries Date: Sun, 6 Apr 2025 17:39:40 +0100 Subject: [PATCH] feat: handle incoming media (images, audio, videos, documents) and show names in group messages (#34) * added support for recieving media + improved group printing * fix formatting --- whatsapp-bridge/main.go | 471 +++++++++++++++++++++++++++----- whatsapp-mcp-server/main.py | 30 +- whatsapp-mcp-server/whatsapp.py | 267 +++++++++--------- 3 files changed, 564 insertions(+), 204 deletions(-) diff --git a/whatsapp-bridge/main.go b/whatsapp-bridge/main.go index 5a4ecd5..d1e0270 100644 --- a/whatsapp-bridge/main.go +++ b/whatsapp-bridge/main.go @@ -11,6 +11,7 @@ import ( "net/http" "os" "os/signal" + "path/filepath" "reflect" "strings" "syscall" @@ -32,10 +33,12 @@ import ( // Message represents a chat message for our client type Message struct { - Time time.Time - Sender string - Content string - IsFromMe bool + Time time.Time + Sender string + Content string + IsFromMe bool + MediaType string + Filename string } // Database handler for storing message history @@ -71,6 +74,13 @@ func NewMessageStore() (*MessageStore, error) { content TEXT, timestamp TIMESTAMP, is_from_me BOOLEAN, + media_type TEXT, + filename TEXT, + url TEXT, + media_key BLOB, + file_sha256 BLOB, + file_enc_sha256 BLOB, + file_length INTEGER, PRIMARY KEY (id, chat_jid), FOREIGN KEY (chat_jid) REFERENCES chats(jid) ); @@ -98,15 +108,18 @@ func (store *MessageStore) StoreChat(jid, name string, lastMessageTime time.Time } // Store a message in the database -func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, timestamp time.Time, isFromMe bool) error { - // Only store if there's actual content - if content == "" { +func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, timestamp time.Time, isFromMe bool, + mediaType, filename, url string, mediaKey, fileSHA256, fileEncSHA256 []byte, fileLength uint64) error { + // Only store if there's actual content or media + if content == "" && mediaType == "" { return nil } _, err := store.db.Exec( - "INSERT OR REPLACE INTO messages (id, chat_jid, sender, content, timestamp, is_from_me) VALUES (?, ?, ?, ?, ?, ?)", - id, chatJID, sender, content, timestamp, isFromMe, + `INSERT OR REPLACE INTO messages + (id, chat_jid, sender, content, timestamp, is_from_me, media_type, filename, url, media_key, file_sha256, file_enc_sha256, file_length) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + id, chatJID, sender, content, timestamp, isFromMe, mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength, ) return err } @@ -114,7 +127,7 @@ func (store *MessageStore) StoreMessage(id, chatJID, sender, content string, tim // Get messages from a chat func (store *MessageStore) GetMessages(chatJID string, limit int) ([]Message, error) { rows, err := store.db.Query( - "SELECT sender, content, timestamp, is_from_me FROM messages WHERE chat_jid = ? ORDER BY timestamp DESC LIMIT ?", + "SELECT sender, content, timestamp, is_from_me, media_type, filename FROM messages WHERE chat_jid = ? ORDER BY timestamp DESC LIMIT ?", chatJID, limit, ) if err != nil { @@ -126,7 +139,7 @@ func (store *MessageStore) GetMessages(chatJID string, limit int) ([]Message, er for rows.Next() { var msg Message var timestamp time.Time - err := rows.Scan(&msg.Sender, &msg.Content, ×tamp, &msg.IsFromMe) + err := rows.Scan(&msg.Sender, &msg.Content, ×tamp, &msg.IsFromMe, &msg.MediaType, &msg.Filename) if err != nil { return nil, err } @@ -358,8 +371,312 @@ func sendWhatsAppMessage(client *whatsmeow.Client, recipient string, message str return true, fmt.Sprintf("Message sent to %s", recipient) } +// Extract media info from a message +func extractMediaInfo(msg *waProto.Message) (mediaType string, filename string, url string, mediaKey []byte, fileSHA256 []byte, fileEncSHA256 []byte, fileLength uint64) { + if msg == nil { + return "", "", "", nil, nil, nil, 0 + } + + // Check for image message + if img := msg.GetImageMessage(); img != nil { + return "image", "image_" + time.Now().Format("20060102_150405") + ".jpg", + img.GetURL(), img.GetMediaKey(), img.GetFileSHA256(), img.GetFileEncSHA256(), img.GetFileLength() + } + + // Check for video message + if vid := msg.GetVideoMessage(); vid != nil { + return "video", "video_" + time.Now().Format("20060102_150405") + ".mp4", + vid.GetURL(), vid.GetMediaKey(), vid.GetFileSHA256(), vid.GetFileEncSHA256(), vid.GetFileLength() + } + + // Check for audio message + if aud := msg.GetAudioMessage(); aud != nil { + return "audio", "audio_" + time.Now().Format("20060102_150405") + ".ogg", + aud.GetURL(), aud.GetMediaKey(), aud.GetFileSHA256(), aud.GetFileEncSHA256(), aud.GetFileLength() + } + + // Check for document message + if doc := msg.GetDocumentMessage(); doc != nil { + filename := doc.GetFileName() + if filename == "" { + filename = "document_" + time.Now().Format("20060102_150405") + } + return "document", filename, + doc.GetURL(), doc.GetMediaKey(), doc.GetFileSHA256(), doc.GetFileEncSHA256(), doc.GetFileLength() + } + + return "", "", "", nil, nil, nil, 0 +} + +// Handle regular incoming messages with media support +func handleMessage(client *whatsmeow.Client, messageStore *MessageStore, msg *events.Message, logger waLog.Logger) { + // Save message to database + chatJID := msg.Info.Chat.String() + sender := msg.Info.Sender.User + + // Get appropriate chat name (pass nil for conversation since we don't have one for regular messages) + name := GetChatName(client, messageStore, msg.Info.Chat, chatJID, nil, sender, logger) + + // Update chat in database with the message timestamp (keeps last message time updated) + err := messageStore.StoreChat(chatJID, name, msg.Info.Timestamp) + if err != nil { + logger.Warnf("Failed to store chat: %v", err) + } + + // Extract text content + content := extractTextContent(msg.Message) + + // Extract media info + mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength := extractMediaInfo(msg.Message) + + // Skip if there's no content and no media + if content == "" && mediaType == "" { + return + } + + // Store message in database + err = messageStore.StoreMessage( + msg.Info.ID, + chatJID, + sender, + content, + msg.Info.Timestamp, + msg.Info.IsFromMe, + mediaType, + filename, + url, + mediaKey, + fileSHA256, + fileEncSHA256, + fileLength, + ) + + if err != nil { + logger.Warnf("Failed to store message: %v", err) + } else { + // Log message reception + timestamp := msg.Info.Timestamp.Format("2006-01-02 15:04:05") + direction := "←" + if msg.Info.IsFromMe { + direction = "→" + } + + // Log based on message type + if mediaType != "" { + fmt.Printf("[%s] %s %s: [%s: %s] %s\n", timestamp, direction, sender, mediaType, filename, content) + } else if content != "" { + fmt.Printf("[%s] %s %s: %s\n", timestamp, direction, sender, content) + } + } +} + +// DownloadMediaRequest represents the request body for the download media API +type DownloadMediaRequest struct { + MessageID string `json:"message_id"` + ChatJID string `json:"chat_jid"` +} + +// DownloadMediaResponse represents the response for the download media API +type DownloadMediaResponse struct { + Success bool `json:"success"` + Message string `json:"message"` + Filename string `json:"filename,omitempty"` + Path string `json:"path,omitempty"` +} + +// Store additional media info in the database +func (store *MessageStore) StoreMediaInfo(id, chatJID, url string, mediaKey, fileSHA256, fileEncSHA256 []byte, fileLength uint64) error { + _, err := store.db.Exec( + "UPDATE messages SET url = ?, media_key = ?, file_sha256 = ?, file_enc_sha256 = ?, file_length = ? WHERE id = ? AND chat_jid = ?", + url, mediaKey, fileSHA256, fileEncSHA256, fileLength, id, chatJID, + ) + return err +} + +// Get media info from the database +func (store *MessageStore) GetMediaInfo(id, chatJID string) (string, string, string, []byte, []byte, []byte, uint64, error) { + var mediaType, filename, url string + var mediaKey, fileSHA256, fileEncSHA256 []byte + var fileLength uint64 + + err := store.db.QueryRow( + "SELECT media_type, filename, url, media_key, file_sha256, file_enc_sha256, file_length FROM messages WHERE id = ? AND chat_jid = ?", + id, chatJID, + ).Scan(&mediaType, &filename, &url, &mediaKey, &fileSHA256, &fileEncSHA256, &fileLength) + + return mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength, err +} + +// MediaDownloader implements the whatsmeow.DownloadableMessage interface +type MediaDownloader struct { + URL string + DirectPath string + MediaKey []byte + FileLength uint64 + FileSHA256 []byte + FileEncSHA256 []byte + MediaType whatsmeow.MediaType +} + +// GetDirectPath implements the DownloadableMessage interface +func (d *MediaDownloader) GetDirectPath() string { + return d.DirectPath +} + +// GetURL implements the DownloadableMessage interface +func (d *MediaDownloader) GetURL() string { + return d.URL +} + +// GetMediaKey implements the DownloadableMessage interface +func (d *MediaDownloader) GetMediaKey() []byte { + return d.MediaKey +} + +// GetFileLength implements the DownloadableMessage interface +func (d *MediaDownloader) GetFileLength() uint64 { + return d.FileLength +} + +// GetFileSHA256 implements the DownloadableMessage interface +func (d *MediaDownloader) GetFileSHA256() []byte { + return d.FileSHA256 +} + +// GetFileEncSHA256 implements the DownloadableMessage interface +func (d *MediaDownloader) GetFileEncSHA256() []byte { + return d.FileEncSHA256 +} + +// GetMediaType implements the DownloadableMessage interface +func (d *MediaDownloader) GetMediaType() whatsmeow.MediaType { + return d.MediaType +} + +// Function to download media from a message +func downloadMedia(client *whatsmeow.Client, messageStore *MessageStore, messageID, chatJID string) (bool, string, string, string, error) { + // Query the database for the message + var mediaType, filename, url string + var mediaKey, fileSHA256, fileEncSHA256 []byte + var fileLength uint64 + var err error + + // First, check if we already have this file + chatDir := fmt.Sprintf("store/%s", strings.ReplaceAll(chatJID, ":", "_")) + localPath := "" + + // Get media info from the database + mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength, err = messageStore.GetMediaInfo(messageID, chatJID) + + if err != nil { + // Try to get basic info if extended info isn't available + err = messageStore.db.QueryRow( + "SELECT media_type, filename FROM messages WHERE id = ? AND chat_jid = ?", + messageID, chatJID, + ).Scan(&mediaType, &filename) + + if err != nil { + return false, "", "", "", fmt.Errorf("failed to find message: %v", err) + } + } + + // Check if this is a media message + if mediaType == "" { + return false, "", "", "", fmt.Errorf("not a media message") + } + + // Create directory for the chat if it doesn't exist + if err := os.MkdirAll(chatDir, 0755); err != nil { + return false, "", "", "", fmt.Errorf("failed to create chat directory: %v", err) + } + + // Generate a local path for the file + localPath = fmt.Sprintf("%s/%s", chatDir, filename) + + // Get absolute path + absPath, err := filepath.Abs(localPath) + if err != nil { + return false, "", "", "", fmt.Errorf("failed to get absolute path: %v", err) + } + + // Check if file already exists + if _, err := os.Stat(localPath); err == nil { + // File exists, return it + return true, mediaType, filename, absPath, nil + } + + // If we don't have all the media info we need, we can't download + if url == "" || len(mediaKey) == 0 || len(fileSHA256) == 0 || len(fileEncSHA256) == 0 || fileLength == 0 { + return false, "", "", "", fmt.Errorf("incomplete media information for download") + } + + fmt.Printf("Attempting to download media for message %s in chat %s...\n", messageID, chatJID) + + // Extract direct path from URL + directPath := extractDirectPathFromURL(url) + + // Create a downloader that implements DownloadableMessage + var waMediaType whatsmeow.MediaType + switch mediaType { + case "image": + waMediaType = whatsmeow.MediaImage + case "video": + waMediaType = whatsmeow.MediaVideo + case "audio": + waMediaType = whatsmeow.MediaAudio + case "document": + waMediaType = whatsmeow.MediaDocument + default: + return false, "", "", "", fmt.Errorf("unsupported media type: %s", mediaType) + } + + downloader := &MediaDownloader{ + URL: url, + DirectPath: directPath, + MediaKey: mediaKey, + FileLength: fileLength, + FileSHA256: fileSHA256, + FileEncSHA256: fileEncSHA256, + MediaType: waMediaType, + } + + // Download the media using whatsmeow client + mediaData, err := client.Download(downloader) + if err != nil { + return false, "", "", "", fmt.Errorf("failed to download media: %v", err) + } + + // Save the downloaded media to file + if err := os.WriteFile(localPath, mediaData, 0644); err != nil { + return false, "", "", "", fmt.Errorf("failed to save media file: %v", err) + } + + fmt.Printf("Successfully downloaded %s media to %s (%d bytes)\n", mediaType, absPath, len(mediaData)) + return true, mediaType, filename, absPath, nil +} + +// Extract direct path from a WhatsApp media URL +func extractDirectPathFromURL(url string) string { + // The direct path is typically in the URL, we need to extract it + // Example URL: https://mmg.whatsapp.net/v/t62.7118-24/13812002_698058036224062_3424455886509161511_n.enc?ccb=11-4&oh=... + + // Find the path part after the domain + parts := strings.SplitN(url, ".net/", 2) + if len(parts) < 2 { + return url // Return original URL if parsing fails + } + + pathPart := parts[1] + + // Remove query parameters + pathPart = strings.SplitN(pathPart, "?", 2)[0] + + // Create proper direct path format + return "/" + pathPart +} + // Start a REST API server to expose the WhatsApp client functionality -func startRESTServer(client *whatsmeow.Client, port int) { +func startRESTServer(client *whatsmeow.Client, messageStore *MessageStore, port int) { // Handler for sending messages http.HandleFunc("/api/send", func(w http.ResponseWriter, r *http.Request) { // Only allow POST requests @@ -406,6 +723,57 @@ func startRESTServer(client *whatsmeow.Client, port int) { }) }) + // Handler for downloading media + http.HandleFunc("/api/download", func(w http.ResponseWriter, r *http.Request) { + // Only allow POST requests + if r.Method != http.MethodPost { + http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) + return + } + + // Parse the request body + var req DownloadMediaRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, "Invalid request format", http.StatusBadRequest) + return + } + + // Validate request + if req.MessageID == "" || req.ChatJID == "" { + http.Error(w, "Message ID and Chat JID are required", http.StatusBadRequest) + return + } + + // Download the media + success, mediaType, filename, path, err := downloadMedia(client, messageStore, req.MessageID, req.ChatJID) + + // Set response headers + w.Header().Set("Content-Type", "application/json") + + // Handle download result + if !success || err != nil { + errMsg := "Unknown error" + if err != nil { + errMsg = err.Error() + } + + w.WriteHeader(http.StatusInternalServerError) + json.NewEncoder(w).Encode(DownloadMediaResponse{ + Success: false, + Message: fmt.Sprintf("Failed to download media: %s", errMsg), + }) + return + } + + // Send successful response + json.NewEncoder(w).Encode(DownloadMediaResponse{ + Success: true, + Message: fmt.Sprintf("Successfully downloaded %s media", mediaType), + Filename: filename, + Path: path, + }) + }) + // Start the server serverAddr := fmt.Sprintf(":%d", port) fmt.Printf("Starting REST API server on %s...\n", serverAddr) @@ -538,7 +906,7 @@ func main() { fmt.Println("\n✓ Connected to WhatsApp! Type 'help' for commands.") // Start REST API server - startRESTServer(client, 8080) + startRESTServer(client, messageStore, 8080) // Create a channel to keep the main goroutine alive exitChan := make(chan os.Signal, 1) @@ -637,49 +1005,6 @@ func GetChatName(client *whatsmeow.Client, messageStore *MessageStore, jid types return name } -// Handle regular incoming messages -func handleMessage(client *whatsmeow.Client, messageStore *MessageStore, msg *events.Message, logger waLog.Logger) { - // Extract text content - content := extractTextContent(msg.Message) - if content == "" { - return // Skip non-text messages - } - - // Save message to database - chatJID := msg.Info.Chat.String() - sender := msg.Info.Sender.User - - // Get appropriate chat name (pass nil for conversation since we don't have one for regular messages) - name := GetChatName(client, messageStore, msg.Info.Chat, chatJID, nil, sender, logger) - - // Update chat in database with the message timestamp (keeps last message time updated) - err := messageStore.StoreChat(chatJID, name, msg.Info.Timestamp) - if err != nil { - logger.Warnf("Failed to store chat: %v", err) - } - - // Store message in database - err = messageStore.StoreMessage( - msg.Info.ID, - chatJID, - sender, - content, - msg.Info.Timestamp, - msg.Info.IsFromMe, - ) - if err != nil { - logger.Warnf("Failed to store message: %v", err) - } else { - // Log message reception - timestamp := msg.Info.Timestamp.Format("2006-01-02 15:04:05") - direction := "←" - if msg.Info.IsFromMe { - direction = "→" - } - fmt.Printf("[%s] %s %s: %s\n", timestamp, direction, sender, content) - } -} - // Handle history sync events func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, historySync *events.HistorySync, logger waLog.Logger) { fmt.Printf("Received history sync event with %d conversations\n", len(historySync.Data.Conversations)) @@ -738,11 +1063,20 @@ func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, his } } - // Log the message content for debugging - logger.Infof("Message content: %v", content) + // Extract media info + var mediaType, filename, url string + var mediaKey, fileSHA256, fileEncSHA256 []byte + var fileLength uint64 - // Skip non-text messages - if content == "" { + if msg.Message.Message != nil { + mediaType, filename, url, mediaKey, fileSHA256, fileEncSHA256, fileLength = extractMediaInfo(msg.Message.Message) + } + + // Log the message content for debugging + logger.Infof("Message content: %v, Media Type: %v", content, mediaType) + + // Skip messages with no content and no media + if content == "" && mediaType == "" { continue } @@ -785,19 +1119,32 @@ func handleHistorySync(client *whatsmeow.Client, messageStore *MessageStore, his content, timestamp, isFromMe, + mediaType, + filename, + url, + mediaKey, + fileSHA256, + fileEncSHA256, + fileLength, ) if err != nil { logger.Warnf("Failed to store history message: %v", err) } else { syncedCount++ // Log successful message storage - logger.Infof("Stored message: [%s] %s -> %s: %s", timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, content) + if mediaType != "" { + logger.Infof("Stored message: [%s] %s -> %s: [%s: %s] %s", + timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, mediaType, filename, content) + } else { + logger.Infof("Stored message: [%s] %s -> %s: %s", + timestamp.Format("2006-01-02 15:04:05"), sender, chatJID, content) + } } } } } - fmt.Printf("History sync complete. Stored %d text messages.\n", syncedCount) + fmt.Printf("History sync complete. Stored %d messages.\n", syncedCount) } // Request history sync from the server diff --git a/whatsapp-mcp-server/main.py b/whatsapp-mcp-server/main.py index 216eaaf..ce1b85b 100644 --- a/whatsapp-mcp-server/main.py +++ b/whatsapp-mcp-server/main.py @@ -11,7 +11,8 @@ from whatsapp import ( get_message_context as whatsapp_get_message_context, send_message as whatsapp_send_message, send_file as whatsapp_send_file, - send_audio_message as whatsapp_audio_voice_message + send_audio_message as whatsapp_audio_voice_message, + download_media as whatsapp_download_media ) # Initialize FastMCP server @@ -128,7 +129,7 @@ def get_contact_chats(jid: str, limit: int = 20, page: int = 0) -> List[Dict[str return chats @mcp.tool() -def get_last_interaction(jid: str) -> Dict[str, Any]: +def get_last_interaction(jid: str) -> str: """Get most recent WhatsApp message involving the contact. Args: @@ -220,6 +221,31 @@ def send_audio_message(recipient: str, media_path: str) -> Dict[str, Any]: "message": status_message } +@mcp.tool() +def download_media(message_id: str, chat_jid: str) -> Dict[str, Any]: + """Download media from a WhatsApp message and get the local file path. + + Args: + message_id: The ID of the message containing the media + chat_jid: The JID of the chat containing the message + + Returns: + A dictionary containing success status, a status message, and the file path if successful + """ + file_path = whatsapp_download_media(message_id, chat_jid) + + if file_path: + return { + "success": True, + "message": "Media downloaded successfully", + "file_path": file_path + } + else: + return { + "success": False, + "message": "Failed to download media" + } + if __name__ == "__main__": # Initialize and run the server mcp.run(transport='stdio') \ No newline at end of file diff --git a/whatsapp-mcp-server/whatsapp.py b/whatsapp-mcp-server/whatsapp.py index cc6f4ee..55ff50a 100644 --- a/whatsapp-mcp-server/whatsapp.py +++ b/whatsapp-mcp-server/whatsapp.py @@ -19,6 +19,7 @@ class Message: chat_jid: str id: str chat_name: Optional[str] = None + media_type: Optional[str] = None @dataclass class Chat: @@ -46,146 +47,80 @@ class MessageContext: before: List[Message] after: List[Message] -def print_message(message: Message, show_chat_info: bool = True) -> None: - """Print a single message with consistent formatting.""" - direction = "→" if message.is_from_me else "←" - - if show_chat_info and message.chat_name: - print(f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] {direction} Chat: {message.chat_name} ({message.chat_jid})") - else: - print(f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] {direction}") - - print(f"From: {'Me' if message.is_from_me else message.sender}") - print(f"Message: {message.content}") - print("-" * 100) - -def print_messages_list(messages: List[Message], title: str = "", show_chat_info: bool = True) -> None: - """Print a list of messages with a title and consistent formatting.""" - if not messages: - print("No messages to display.") - return - - if title: - print(f"\n{title}") - print("-" * 100) - - for message in messages: - print_message(message, show_chat_info) - -def print_chat(chat: Chat) -> None: - """Print a single chat with consistent formatting.""" - print(f"Chat: {chat.name} ({chat.jid})") - if chat.last_message_time: - print(f"Last active: {chat.last_message_time:%Y-%m-%d %H:%M:%S}") - direction = "→" if chat.last_is_from_me else "←" - sender = "Me" if chat.last_is_from_me else chat.last_sender - print(f"Last message: {direction} {sender}: {chat.last_message}") - print("-" * 100) - -def print_chats_list(chats: List[Chat], title: str = "") -> None: - """Print a list of chats with a title and consistent formatting.""" - if not chats: - print("No chats to display.") - return - - if title: - print(f"\n{title}") - print("-" * 100) - - for chat in chats: - print_chat(chat) - -def print_paginated_messages(messages: List[Message], page: int, total_pages: int, chat_name: str) -> None: - """Print a paginated list of messages with navigation hints.""" - print(f"\nMessages for chat: {chat_name}") - print(f"Page {page} of {total_pages}") - print("-" * 100) - - print_messages_list(messages, show_chat_info=False) - - # Print pagination info - if page > 1: - print(f"Use page={page-1} to see newer messages") - if page < total_pages: - print(f"Use page={page+1} to see older messages") - -""" -CREATE TABLE messages ( - id TEXT, - chat_jid TEXT, - sender TEXT, - content TEXT, - timestamp TIMESTAMP, - is_from_me BOOLEAN, - PRIMARY KEY (id, chat_jid), - FOREIGN KEY (chat_jid) REFERENCES chats(jid) - ) - -CREATE TABLE chats ( - jid TEXT PRIMARY KEY, - name TEXT, - last_message_time TIMESTAMP - ) -""" - -def print_recent_messages(limit=10) -> List[Message]: +def get_sender_name(sender_jid: str) -> str: try: - # Connect to the SQLite database conn = sqlite3.connect(MESSAGES_DB_PATH) cursor = conn.cursor() - # Query recent messages with chat info - query = """ - SELECT - m.timestamp, - m.sender, - c.name, - m.content, - m.is_from_me, - c.jid, - m.id - FROM messages m - JOIN chats c ON m.chat_jid = c.jid - ORDER BY m.timestamp DESC - LIMIT ? - """ + # First try matching by exact JID + cursor.execute(""" + SELECT name + FROM chats + WHERE jid = ? + LIMIT 1 + """, (sender_jid,)) - cursor.execute(query, (limit,)) - messages = cursor.fetchall() + result = cursor.fetchone() - if not messages: - print("No messages found in the database.") - return [] + # If no result, try looking for the number within JIDs + if not result: + # Extract the phone number part if it's a JID + if '@' in sender_jid: + phone_part = sender_jid.split('@')[0] + else: + phone_part = sender_jid + + cursor.execute(""" + SELECT name + FROM chats + WHERE jid LIKE ? + LIMIT 1 + """, (f"%{phone_part}%",)) - result = [] + result = cursor.fetchone() - # Convert to Message objects - for msg in messages: - message = Message( - timestamp=datetime.fromisoformat(msg[0]), - sender=msg[1], - chat_name=msg[2] or "Unknown Chat", - content=msg[3], - is_from_me=msg[4], - chat_jid=msg[5], - id=msg[6] - ) - result.append(message) + if result and result[0]: + return result[0] + else: + return sender_jid - # Print messages using helper function - print_messages_list(result, title=f"Last {limit} messages:") - return result - except sqlite3.Error as e: - print(f"Error accessing database: {e}") - return [] - except Exception as e: - print(f"Unexpected error: {e}") - return [] + print(f"Database error while getting sender name: {e}") + return sender_jid finally: if 'conn' in locals(): conn.close() +def format_message(message: Message, show_chat_info: bool = True) -> None: + """Print a single message with consistent formatting.""" + output = "" + + if show_chat_info and message.chat_name: + output += f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] Chat: {message.chat_name} " + else: + output += f"[{message.timestamp:%Y-%m-%d %H:%M:%S}] " + + content_prefix = "" + if hasattr(message, 'media_type') and message.media_type: + content_prefix = f"[{message.media_type} - Message ID: {message.id} - Chat JID: {message.chat_jid}] " + + try: + sender_name = get_sender_name(message.sender) if not message.is_from_me else "Me" + output += f"From: {sender_name}: {content_prefix}{message.content}\n" + except Exception as e: + print(f"Error formatting message: {e}") + return output + +def format_messages_list(messages: List[Message], show_chat_info: bool = True) -> None: + output = "" + if not messages: + output += "No messages to display." + return output + + for message in messages: + output += format_message(message, show_chat_info) + return output + def list_messages( after: Optional[str] = None, before: Optional[str] = None, @@ -204,7 +139,7 @@ def list_messages( cursor = conn.cursor() # Build base query - query_parts = ["SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id FROM messages"] + query_parts = ["SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type FROM messages"] query_parts.append("JOIN chats ON messages.chat_jid = chats.jid") where_clauses = [] params = [] @@ -261,7 +196,8 @@ def list_messages( content=msg[3], is_from_me=msg[4], chat_jid=msg[5], - id=msg[6] + id=msg[6], + media_type=msg[7] ) result.append(message) @@ -273,9 +209,11 @@ def list_messages( messages_with_context.extend(context.before) messages_with_context.append(context.message) messages_with_context.extend(context.after) - return messages_with_context - return result + return format_messages_list(messages_with_context, show_chat_info=True) + + # Format and display messages without context + return format_messages_list(result, show_chat_info=True) except sqlite3.Error as e: print(f"Database error: {e}") @@ -297,7 +235,7 @@ def get_message_context( # Get the target message first cursor.execute(""" - SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.chat_jid + SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.chat_jid, messages.media_type FROM messages JOIN chats ON messages.chat_jid = chats.jid WHERE messages.id = ? @@ -314,12 +252,13 @@ def get_message_context( content=msg_data[3], is_from_me=msg_data[4], chat_jid=msg_data[5], - id=msg_data[6] + id=msg_data[6], + media_type=msg_data[8] ) # Get messages before cursor.execute(""" - SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id + SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type FROM messages JOIN chats ON messages.chat_jid = chats.jid WHERE messages.chat_jid = ? AND messages.timestamp < ? @@ -336,12 +275,13 @@ def get_message_context( content=msg[3], is_from_me=msg[4], chat_jid=msg[5], - id=msg[6] + id=msg[6], + media_type=msg[7] )) # Get messages after cursor.execute(""" - SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id + SELECT messages.timestamp, messages.sender, chats.name, messages.content, messages.is_from_me, chats.jid, messages.id, messages.media_type FROM messages JOIN chats ON messages.chat_jid = chats.jid WHERE messages.chat_jid = ? AND messages.timestamp > ? @@ -358,7 +298,8 @@ def get_message_context( content=msg[3], is_from_me=msg[4], chat_jid=msg[5], - id=msg[6] + id=msg[6], + media_type=msg[7] )) return MessageContext( @@ -542,7 +483,7 @@ def get_contact_chats(jid: str, limit: int = 20, page: int = 0) -> List[Chat]: conn.close() -def get_last_interaction(jid: str) -> Optional[Message]: +def get_last_interaction(jid: str) -> str: """Get most recent message involving the contact.""" try: conn = sqlite3.connect(MESSAGES_DB_PATH) @@ -556,7 +497,8 @@ def get_last_interaction(jid: str) -> Optional[Message]: m.content, m.is_from_me, c.jid, - m.id + m.id, + m.media_type FROM messages m JOIN chats c ON m.chat_jid = c.jid WHERE m.sender = ? OR c.jid = ? @@ -569,16 +511,19 @@ def get_last_interaction(jid: str) -> Optional[Message]: if not msg_data: return None - return Message( + message = Message( timestamp=datetime.fromisoformat(msg_data[0]), sender=msg_data[1], chat_name=msg_data[2], content=msg_data[3], is_from_me=msg_data[4], chat_jid=msg_data[5], - id=msg_data[6] + id=msg_data[6], + media_type=msg_data[7] ) + return format_message(message) + except sqlite3.Error as e: print(f"Database error: {e}") return None @@ -778,3 +723,45 @@ def send_audio_message(recipient: str, media_path: str) -> Tuple[bool, str]: return False, f"Error parsing response: {response.text}" except Exception as e: return False, f"Unexpected error: {str(e)}" + +def download_media(message_id: str, chat_jid: str) -> Optional[str]: + """Download media from a message and return the local file path. + + Args: + message_id: The ID of the message containing the media + chat_jid: The JID of the chat containing the message + + Returns: + The local file path if download was successful, None otherwise + """ + try: + url = f"{WHATSAPP_API_BASE_URL}/download" + payload = { + "message_id": message_id, + "chat_jid": chat_jid + } + + response = requests.post(url, json=payload) + + if response.status_code == 200: + result = response.json() + if result.get("success", False): + path = result.get("path") + print(f"Media downloaded successfully: {path}") + return path + else: + print(f"Download failed: {result.get('message', 'Unknown error')}") + return None + else: + print(f"Error: HTTP {response.status_code} - {response.text}") + return None + + except requests.RequestException as e: + print(f"Request error: {str(e)}") + return None + except json.JSONDecodeError: + print(f"Error parsing response: {response.text}") + return None + except Exception as e: + print(f"Unexpected error: {str(e)}") + return None