portal.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2021 Tulir Asokan
  3. //
  4. // This program is free software: you can redistribute it and/or modify
  5. // it under the terms of the GNU Affero General Public License as published by
  6. // the Free Software Foundation, either version 3 of the License, or
  7. // (at your option) any later version.
  8. //
  9. // This program is distributed in the hope that it will be useful,
  10. // but WITHOUT ANY WARRANTY; without even the implied warranty of
  11. // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
  12. // GNU Affero General Public License for more details.
  13. //
  14. // You should have received a copy of the GNU Affero General Public License
  15. // along with this program. If not, see <https://www.gnu.org/licenses/>.
  16. package database
  17. import (
  18. "database/sql"
  19. log "maunium.net/go/maulogger/v2"
  20. "maunium.net/go/mautrix/id"
  21. "go.mau.fi/whatsmeow/types"
  22. )
  23. type PortalKey struct {
  24. JID types.JID
  25. Receiver types.JID
  26. }
  27. func GroupPortalKey(jid types.JID) PortalKey {
  28. return PortalKey{
  29. JID: jid.ToNonAD(),
  30. Receiver: jid.ToNonAD(),
  31. }
  32. }
  33. func NewPortalKey(jid, receiver types.JID) PortalKey {
  34. if jid.Server == types.GroupServer {
  35. receiver = jid
  36. }
  37. return PortalKey{
  38. JID: jid.ToNonAD(),
  39. Receiver: receiver.ToNonAD(),
  40. }
  41. }
  42. func (key PortalKey) String() string {
  43. if key.Receiver == key.JID {
  44. return key.JID.String()
  45. }
  46. return key.JID.String() + "-" + key.Receiver.String()
  47. }
  48. type PortalQuery struct {
  49. db *Database
  50. log log.Logger
  51. }
  52. func (pq *PortalQuery) New() *Portal {
  53. return &Portal{
  54. db: pq.db,
  55. log: pq.log,
  56. }
  57. }
  58. func (pq *PortalQuery) GetAll() []*Portal {
  59. return pq.getAll("SELECT * FROM portal")
  60. }
  61. func (pq *PortalQuery) GetByJID(key PortalKey) *Portal {
  62. return pq.get("SELECT * FROM portal WHERE jid=$1 AND receiver=$2", key.JID, key.Receiver)
  63. }
  64. func (pq *PortalQuery) GetByMXID(mxid id.RoomID) *Portal {
  65. return pq.get("SELECT * FROM portal WHERE mxid=$1", mxid)
  66. }
  67. func (pq *PortalQuery) GetAllByJID(jid types.JID) []*Portal {
  68. return pq.getAll("SELECT * FROM portal WHERE jid=$1", jid)
  69. }
  70. func (pq *PortalQuery) FindPrivateChats(receiver types.JID) []*Portal {
  71. return pq.getAll("SELECT * FROM portal WHERE receiver='$1@s.whatsapp.net' AND jid LIKE '%@s.whatsapp.net'", receiver)
  72. }
  73. func (pq *PortalQuery) getAll(query string, args ...interface{}) (portals []*Portal) {
  74. rows, err := pq.db.Query(query, args...)
  75. if err != nil || rows == nil {
  76. return nil
  77. }
  78. defer rows.Close()
  79. for rows.Next() {
  80. portals = append(portals, pq.New().Scan(rows))
  81. }
  82. return
  83. }
  84. func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
  85. row := pq.db.QueryRow(query, args...)
  86. if row == nil {
  87. return nil
  88. }
  89. return pq.New().Scan(row)
  90. }
  91. type Portal struct {
  92. db *Database
  93. log log.Logger
  94. Key PortalKey
  95. MXID id.RoomID
  96. Name string
  97. Topic string
  98. Avatar string
  99. AvatarURL id.ContentURI
  100. Encrypted bool
  101. }
  102. func (portal *Portal) Scan(row Scannable) *Portal {
  103. var mxid, avatarURL sql.NullString
  104. err := row.Scan(&portal.Key.JID, &portal.Key.Receiver, &mxid, &portal.Name, &portal.Topic, &portal.Avatar, &avatarURL, &portal.Encrypted)
  105. if err != nil {
  106. if err != sql.ErrNoRows {
  107. portal.log.Errorln("Database scan failed:", err)
  108. }
  109. return nil
  110. }
  111. portal.MXID = id.RoomID(mxid.String)
  112. portal.AvatarURL, _ = id.ParseContentURI(avatarURL.String)
  113. return portal
  114. }
  115. func (portal *Portal) mxidPtr() *id.RoomID {
  116. if len(portal.MXID) > 0 {
  117. return &portal.MXID
  118. }
  119. return nil
  120. }
  121. func (portal *Portal) Insert() {
  122. _, err := portal.db.Exec("INSERT INTO portal (jid, receiver, mxid, name, topic, avatar, avatar_url, encrypted) VALUES ($1, $2, $3, $4, $5, $6, $7, $8)",
  123. portal.Key.JID, portal.Key.Receiver, portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted)
  124. if err != nil {
  125. portal.log.Warnfln("Failed to insert %s: %v", portal.Key, err)
  126. }
  127. }
  128. func (portal *Portal) Update() {
  129. var mxid *id.RoomID
  130. if len(portal.MXID) > 0 {
  131. mxid = &portal.MXID
  132. }
  133. _, err := portal.db.Exec("UPDATE portal SET mxid=$1, name=$2, topic=$3, avatar=$4, avatar_url=$5, encrypted=$6 WHERE jid=$7 AND receiver=$8",
  134. mxid, portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted, portal.Key.JID, portal.Key.Receiver)
  135. if err != nil {
  136. portal.log.Warnfln("Failed to update %s: %v", portal.Key, err)
  137. }
  138. }
  139. func (portal *Portal) Delete() {
  140. _, err := portal.db.Exec("DELETE FROM portal WHERE jid=$1 AND receiver=$2", portal.Key.JID, portal.Key.Receiver)
  141. if err != nil {
  142. portal.log.Warnfln("Failed to delete %s: %v", portal.Key, err)
  143. }
  144. }
  145. //func (portal *Portal) GetUserIDs() []id.UserID {
  146. // rows, err := portal.db.Query(`SELECT "user".mxid FROM "user", user_portal
  147. // WHERE "user".jid=user_portal.user_jid
  148. // AND user_portal.portal_jid=$1
  149. // AND user_portal.portal_receiver=$2`,
  150. // portal.Key.JID, portal.Key.Receiver)
  151. // if err != nil {
  152. // portal.log.Debugln("Failed to get portal user ids:", err)
  153. // return nil
  154. // }
  155. // var userIDs []id.UserID
  156. // for rows.Next() {
  157. // var userID id.UserID
  158. // err = rows.Scan(&userID)
  159. // if err != nil {
  160. // portal.log.Warnln("Failed to scan row:", err)
  161. // continue
  162. // }
  163. // userIDs = append(userIDs, userID)
  164. // }
  165. // return userIDs
  166. //}