user.go 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247
  1. // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
  2. // Copyright (C) 2020 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. "fmt"
  20. "strings"
  21. "time"
  22. "github.com/Rhymen/go-whatsapp"
  23. log "maunium.net/go/maulogger/v2"
  24. "maunium.net/go/mautrix/id"
  25. )
  26. type UserQuery struct {
  27. db *Database
  28. log log.Logger
  29. }
  30. func (uq *UserQuery) New() *User {
  31. return &User{
  32. db: uq.db,
  33. log: uq.log,
  34. }
  35. }
  36. func (uq *UserQuery) GetAll() (users []*User) {
  37. rows, err := uq.db.Query(`SELECT mxid, jid, management_room, last_connection, client_id, client_token, server_token, enc_key, mac_key FROM "user"`)
  38. if err != nil || rows == nil {
  39. return nil
  40. }
  41. defer rows.Close()
  42. for rows.Next() {
  43. users = append(users, uq.New().Scan(rows))
  44. }
  45. return
  46. }
  47. func (uq *UserQuery) GetByMXID(userID id.UserID) *User {
  48. row := uq.db.QueryRow(`SELECT mxid, jid, management_room, last_connection, client_id, client_token, server_token, enc_key, mac_key FROM "user" WHERE mxid=$1`, userID)
  49. if row == nil {
  50. return nil
  51. }
  52. return uq.New().Scan(row)
  53. }
  54. func (uq *UserQuery) GetByJID(userID whatsapp.JID) *User {
  55. row := uq.db.QueryRow(`SELECT mxid, jid, management_room, last_connection, client_id, client_token, server_token, enc_key, mac_key FROM "user" WHERE jid=$1`, stripSuffix(userID))
  56. if row == nil {
  57. return nil
  58. }
  59. return uq.New().Scan(row)
  60. }
  61. type User struct {
  62. db *Database
  63. log log.Logger
  64. MXID id.UserID
  65. JID whatsapp.JID
  66. ManagementRoom id.RoomID
  67. Session *whatsapp.Session
  68. LastConnection uint64
  69. }
  70. func (user *User) Scan(row Scannable) *User {
  71. var jid, clientID, clientToken, serverToken sql.NullString
  72. var encKey, macKey []byte
  73. err := row.Scan(&user.MXID, &jid, &user.ManagementRoom, &user.LastConnection, &clientID, &clientToken, &serverToken, &encKey, &macKey)
  74. if err != nil {
  75. if err != sql.ErrNoRows {
  76. user.log.Errorln("Database scan failed:", err)
  77. }
  78. return nil
  79. }
  80. if len(jid.String) > 0 && len(clientID.String) > 0 {
  81. user.JID = jid.String + whatsapp.NewUserSuffix
  82. user.Session = &whatsapp.Session{
  83. ClientID: clientID.String,
  84. ClientToken: clientToken.String,
  85. ServerToken: serverToken.String,
  86. EncKey: encKey,
  87. MacKey: macKey,
  88. Wid: jid.String + whatsapp.OldUserSuffix,
  89. }
  90. } else {
  91. user.Session = nil
  92. }
  93. return user
  94. }
  95. func stripSuffix(jid whatsapp.JID) string {
  96. if len(jid) == 0 {
  97. return jid
  98. }
  99. index := strings.IndexRune(jid, '@')
  100. if index < 0 {
  101. return jid
  102. }
  103. return jid[:index]
  104. }
  105. func (user *User) jidPtr() *string {
  106. if len(user.JID) > 0 {
  107. str := stripSuffix(user.JID)
  108. return &str
  109. }
  110. return nil
  111. }
  112. func (user *User) sessionUnptr() (sess whatsapp.Session) {
  113. if user.Session != nil {
  114. sess = *user.Session
  115. }
  116. return
  117. }
  118. func (user *User) Insert() {
  119. sess := user.sessionUnptr()
  120. _, err := user.db.Exec(`INSERT INTO "user" (mxid, jid, management_room, last_connection, client_id, client_token, server_token, enc_key, mac_key) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)`,
  121. user.MXID, user.jidPtr(),
  122. user.ManagementRoom, user.LastConnection,
  123. sess.ClientID, sess.ClientToken, sess.ServerToken, sess.EncKey, sess.MacKey)
  124. if err != nil {
  125. user.log.Warnfln("Failed to insert %s: %v", user.MXID, err)
  126. }
  127. }
  128. func (user *User) UpdateLastConnection() {
  129. user.LastConnection = uint64(time.Now().Unix())
  130. _, err := user.db.Exec(`UPDATE "user" SET last_connection=$1 WHERE mxid=$2`,
  131. user.LastConnection, user.MXID)
  132. if err != nil {
  133. user.log.Warnfln("Failed to update last connection ts: %v", err)
  134. }
  135. }
  136. func (user *User) Update() {
  137. sess := user.sessionUnptr()
  138. _, err := user.db.Exec(`UPDATE "user" SET jid=$1, management_room=$2, last_connection=$3, client_id=$4, client_token=$5, server_token=$6, enc_key=$7, mac_key=$8 WHERE mxid=$9`,
  139. user.jidPtr(), user.ManagementRoom, user.LastConnection,
  140. sess.ClientID, sess.ClientToken, sess.ServerToken, sess.EncKey, sess.MacKey,
  141. user.MXID)
  142. if err != nil {
  143. user.log.Warnfln("Failed to update %s: %v", user.MXID, err)
  144. }
  145. }
  146. type PortalKeyWithMeta struct {
  147. PortalKey
  148. InCommunity bool
  149. }
  150. func (user *User) SetPortalKeys(newKeys []PortalKeyWithMeta) error {
  151. tx, err := user.db.Begin()
  152. if err != nil {
  153. return err
  154. }
  155. _, err = tx.Exec("DELETE FROM user_portal WHERE user_jid=$1", user.jidPtr())
  156. if err != nil {
  157. _ = tx.Rollback()
  158. return err
  159. }
  160. valueStrings := make([]string, len(newKeys))
  161. values := make([]interface{}, len(newKeys)*4)
  162. for i, key := range newKeys {
  163. pos := i * 4
  164. valueStrings[i] = fmt.Sprintf("($%d, $%d, $%d, $%d)", pos+1, pos+2, pos+3, pos+4)
  165. values[pos] = user.jidPtr()
  166. values[pos+1] = key.JID
  167. values[pos+2] = key.Receiver
  168. values[pos+3] = key.InCommunity
  169. }
  170. query := fmt.Sprintf("INSERT INTO user_portal (user_jid, portal_jid, portal_receiver, in_community) VALUES %s",
  171. strings.Join(valueStrings, ", "))
  172. _, err = tx.Exec(query, values...)
  173. if err != nil {
  174. _ = tx.Rollback()
  175. return err
  176. }
  177. return tx.Commit()
  178. }
  179. func (user *User) IsInPortal(key PortalKey) bool {
  180. row := user.db.QueryRow(`SELECT EXISTS(SELECT 1 FROM user_portal WHERE user_jid=$1 AND portal_jid=$2 AND portal_receiver=$3)`, user.jidPtr(), &key.JID, &key.Receiver)
  181. var exists bool
  182. _ = row.Scan(&exists)
  183. return exists
  184. }
  185. func (user *User) GetPortalKeys() []PortalKey {
  186. rows, err := user.db.Query(`SELECT portal_jid, portal_receiver FROM user_portal WHERE user_jid=$1`, user.jidPtr())
  187. if err != nil {
  188. user.log.Warnln("Failed to get user portal keys:", err)
  189. return nil
  190. }
  191. var keys []PortalKey
  192. for rows.Next() {
  193. var key PortalKey
  194. err = rows.Scan(&key.JID, &key.Receiver)
  195. if err != nil {
  196. user.log.Warnln("Failed to scan row:", err)
  197. continue
  198. }
  199. keys = append(keys, key)
  200. }
  201. return keys
  202. }
  203. func (user *User) GetInCommunityMap() map[PortalKey]bool {
  204. rows, err := user.db.Query(`SELECT portal_jid, portal_receiver, in_community FROM user_portal WHERE user_jid=$1`, user.jidPtr())
  205. if err != nil {
  206. user.log.Warnln("Failed to get user portal keys:", err)
  207. return nil
  208. }
  209. keys := make(map[PortalKey]bool)
  210. for rows.Next() {
  211. var key PortalKey
  212. var inCommunity bool
  213. err = rows.Scan(&key.JID, &key.Receiver, &inCommunity)
  214. if err != nil {
  215. user.log.Warnln("Failed to scan row:", err)
  216. continue
  217. }
  218. keys[key] = inCommunity
  219. }
  220. return keys
  221. }