portal.go 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176
  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 NewPortalKey(jid, receiver types.JID) PortalKey {
  28. if jid.Server == types.GroupServer {
  29. receiver = jid
  30. } else if jid.Server == types.LegacyUserServer {
  31. jid.Server = types.DefaultUserServer
  32. }
  33. return PortalKey{
  34. JID: jid.ToNonAD(),
  35. Receiver: receiver.ToNonAD(),
  36. }
  37. }
  38. func (key PortalKey) String() string {
  39. if key.Receiver == key.JID {
  40. return key.JID.String()
  41. }
  42. return key.JID.String() + "-" + key.Receiver.String()
  43. }
  44. type PortalQuery struct {
  45. db *Database
  46. log log.Logger
  47. }
  48. func (pq *PortalQuery) New() *Portal {
  49. return &Portal{
  50. db: pq.db,
  51. log: pq.log,
  52. }
  53. }
  54. func (pq *PortalQuery) GetAll() []*Portal {
  55. return pq.getAll("SELECT * FROM portal")
  56. }
  57. func (pq *PortalQuery) GetByJID(key PortalKey) *Portal {
  58. return pq.get("SELECT * FROM portal WHERE jid=$1 AND receiver=$2", key.JID, key.Receiver)
  59. }
  60. func (pq *PortalQuery) GetByMXID(mxid id.RoomID) *Portal {
  61. return pq.get("SELECT * FROM portal WHERE mxid=$1", mxid)
  62. }
  63. func (pq *PortalQuery) GetAllByJID(jid types.JID) []*Portal {
  64. return pq.getAll("SELECT * FROM portal WHERE jid=$1", jid.ToNonAD())
  65. }
  66. func (pq *PortalQuery) FindPrivateChats(receiver types.JID) []*Portal {
  67. return pq.getAll("SELECT * FROM portal WHERE receiver=$1 AND jid LIKE '%@s.whatsapp.net'", receiver.ToNonAD())
  68. }
  69. func (pq *PortalQuery) getAll(query string, args ...interface{}) (portals []*Portal) {
  70. rows, err := pq.db.Query(query, args...)
  71. if err != nil || rows == nil {
  72. return nil
  73. }
  74. defer rows.Close()
  75. for rows.Next() {
  76. portals = append(portals, pq.New().Scan(rows))
  77. }
  78. return
  79. }
  80. func (pq *PortalQuery) get(query string, args ...interface{}) *Portal {
  81. row := pq.db.QueryRow(query, args...)
  82. if row == nil {
  83. return nil
  84. }
  85. return pq.New().Scan(row)
  86. }
  87. type Portal struct {
  88. db *Database
  89. log log.Logger
  90. Key PortalKey
  91. MXID id.RoomID
  92. Name string
  93. Topic string
  94. Avatar string
  95. AvatarURL id.ContentURI
  96. Encrypted bool
  97. FirstEventID id.EventID
  98. NextBatchID id.BatchID
  99. RelayUserID id.UserID
  100. }
  101. func (portal *Portal) Scan(row Scannable) *Portal {
  102. var mxid, avatarURL, firstEventID, nextBatchID, relayUserID sql.NullString
  103. err := row.Scan(&portal.Key.JID, &portal.Key.Receiver, &mxid, &portal.Name, &portal.Topic, &portal.Avatar, &avatarURL, &portal.Encrypted, &firstEventID, &nextBatchID, &relayUserID)
  104. if err != nil {
  105. if err != sql.ErrNoRows {
  106. portal.log.Errorln("Database scan failed:", err)
  107. }
  108. return nil
  109. }
  110. portal.MXID = id.RoomID(mxid.String)
  111. portal.AvatarURL, _ = id.ParseContentURI(avatarURL.String)
  112. portal.FirstEventID = id.EventID(firstEventID.String)
  113. portal.NextBatchID = id.BatchID(nextBatchID.String)
  114. portal.RelayUserID = id.UserID(relayUserID.String)
  115. return portal
  116. }
  117. func (portal *Portal) mxidPtr() *id.RoomID {
  118. if len(portal.MXID) > 0 {
  119. return &portal.MXID
  120. }
  121. return nil
  122. }
  123. func (portal *Portal) relayUserPtr() *id.UserID {
  124. if len(portal.RelayUserID) > 0 {
  125. return &portal.RelayUserID
  126. }
  127. return nil
  128. }
  129. func (portal *Portal) Insert() {
  130. _, err := portal.db.Exec("INSERT INTO portal (jid, receiver, mxid, name, topic, avatar, avatar_url, encrypted, first_event_id, next_batch_id, relay_user_id) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)",
  131. portal.Key.JID, portal.Key.Receiver, portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted, portal.FirstEventID.String(), portal.NextBatchID.String(), portal.relayUserPtr())
  132. if err != nil {
  133. portal.log.Warnfln("Failed to insert %s: %v", portal.Key, err)
  134. }
  135. }
  136. func (portal *Portal) Update() {
  137. _, err := portal.db.Exec("UPDATE portal SET mxid=$1, name=$2, topic=$3, avatar=$4, avatar_url=$5, encrypted=$6, first_event_id=$7, next_batch_id=$8, relay_user_id=$9 WHERE jid=$10 AND receiver=$11",
  138. portal.mxidPtr(), portal.Name, portal.Topic, portal.Avatar, portal.AvatarURL.String(), portal.Encrypted, portal.FirstEventID.String(), portal.NextBatchID.String(), portal.relayUserPtr(), portal.Key.JID, portal.Key.Receiver)
  139. if err != nil {
  140. portal.log.Warnfln("Failed to update %s: %v", portal.Key, err)
  141. }
  142. }
  143. func (portal *Portal) Delete() {
  144. _, err := portal.db.Exec("DELETE FROM portal WHERE jid=$1 AND receiver=$2", portal.Key.JID, portal.Key.Receiver)
  145. if err != nil {
  146. portal.log.Warnfln("Failed to delete %s: %v", portal.Key, err)
  147. }
  148. }