Solve Man-in-the-middle attack

bot/src/commands/commands.ts

@@ -20,16 +20,20 @@ export const handleCommands = () => {
       web_app: { url: WEB_APP_URL },
       type: 'web_app',
     })
+
     const query = ctx.update.message.text.split(' ')
     if (query.length === 2) {
-      const username = query[1]
-      return ctx.reply(
-        `الان داری به کاربر ${username} پیام می‌فرستی`,
-        Markup.inlineKeyboard([
-          Markup.button.webApp('ورود به ربات', `${WEB_APP_URL}/@${username}`),
-        ]),
-      )
+      const usernameWithHash = query[1].split('-')
+      if (usernameWithHash.length === 2) {
+        return ctx.reply(
+          `الان داری به کاربر ${usernameWithHash[0]} پیام می‌فرستی`,
+          Markup.inlineKeyboard([
+            Markup.button.webApp('ورود به ربات', `${WEB_APP_URL}/@${query[1]}`),
+          ]),
+        )
+      }
     }
+
     ctx.reply(
       START_MESSAGE,
       Markup.inlineKeyboard([

package.json

@@ -12,7 +12,9 @@
   },
   "dependencies": {
     "@twa-dev/types": "^7.8.0",
+    "@types/crypto-js": "^4.2.2",
     "axios": "^1.6.8",
+    "crypto-js": "^4.2.0",
     "node-forge": "^1.3.1",
     "pinia": "^2.1.7",
     "telegraf": "^4.16.3",

server/database/connection.go

@@ -28,4 +28,5 @@ func InitConnection() error {
 func migration() {
 	DB.AutoMigrate(&User{})
 	DB.AutoMigrate(&Message{})
+	DB.AutoMigrate(&Session{})
 }

server/handlers/handlers.go

@@ -114,7 +114,7 @@ func SetPublicKey(c *fiber.Ctx) error {
 		Model(&result).
 		Where("userid = ?", userid).
 		Update("public_key", body.PublicKey).
-		Update("public_key_hash", utils.GenerateSHA256(body.PublicKey))
+		Update("public_key_hash", utils.GenerateMD5(body.PublicKey))
 
 	return c.JSON(response.SetPublicKey{
 		PublicKey: body.PublicKey,
@@ -192,7 +192,7 @@ func SendMessage(c *fiber.Ctx) error {
 		log.Fatalf("Failed to publish message: %v", err)
 	}
 
-	return c.JSON(response.SendMessage{Message: "The message was sent"})
+	return c.JSON(response.SendMessage{Message: "The message was sent", SessionID: session.ID})
 }
 
 func GetPublicKey(c *fiber.Ctx) error {
@@ -227,12 +227,13 @@ func ReplayMessage(c *fiber.Ctx) error {
 	database.DB.Where("userid = ?", userid).Find(&user)
 
 	database.DB.Create(&database.Message{
-		Content:  body.Message,
-		FromID:   user.ID,
-		ToID:     result.FromID,
-		OwnerID:  result.OwnerID,
-		ParentID: result.ID,
-		Time:     time.Now()})
+		Content:   body.Message,
+		FromID:    user.ID,
+		ToID:      result.FromID,
+		SessionID: result.SessionID,
+		OwnerID:   result.OwnerID,
+		ParentID:  result.ID,
+		Time:      time.Now()})
 
 	var targetUser database.User
 	database.DB.Where("id = ?", result.FromID).Find(&targetUser)
@@ -261,6 +262,7 @@ func GetMessages(c *fiber.Ctx) error {
 
 		var sourceUser database.User
 		database.DB.Where("id = ?", result[i].FromID).Find(&sourceUser)
+
 		if result[i].ParentID != 0 {
 			var res database.Message
 			database.DB.Where("id = ?", result[i].ParentID).Find(&res)
@@ -271,6 +273,7 @@ func GetMessages(c *fiber.Ctx) error {
 					ID:        result[i].ID,
 					Content:   result[i].Content,
 					Time:      result[i].Time,
+					SessionID: result[i].SessionID,
 					Owner:     owner, // true
 					CanReplay: true,
 					Quote: &response.Quote{
@@ -297,16 +300,17 @@ func GetMessages(c *fiber.Ctx) error {
 			database.DB.Where("id = ?", result[i].SessionID).Find(&session)
 
 			messages = append(messages, response.GetMessages{
-				ID:        result[i].ID,
-				SessionID: result[i].SessionID,
-				Time:      result[i].Time,
-				Owner:     owner,
-				Quote:     nil,
-				Content:   result[i].Content,
-				CanReplay: true,
+				ID:         result[i].ID,
+				SessionID:  result[i].SessionID,
+				SessionKey: session.Key,
+				Time:       result[i].Time,
+				Owner:      owner,
+				Quote:      nil,
+				Content:    result[i].Content,
+				CanReplay:  true,
 			})
 
-			database.DB.Delete(&database.Session{}, result[i].SessionID)
+			// database.DB.Delete(&database.Session{}, result[i].SessionID)
 		}
 
 	}

server/request/request.go

@@ -9,7 +9,7 @@ type SetPublicKey struct {
 }
 
 type SendMessage struct {
-	Id         string `json:"id"`
+	Id         uint64 `json:"id"`
 	Message    string `json:"message"`
 	SessionKey string `json:"session_key"`
 }

server/response/response.go

@@ -32,7 +32,8 @@ type GetProfile struct {
 }
 
 type SendMessage struct {
-	Message string `json:"message"`
+	Message   string `json:"message"`
+	SessionID uint64 `json:"session_id"`
 }
 
 type GetMe struct {

server/utils/utils.go

@@ -1,7 +1,7 @@
 package utils
 
 import (
-	"crypto/sha256"
+	"crypto/md5"
 	"encoding/hex"
 	"encoding/json"
 	"errors"
@@ -76,8 +76,8 @@ func Parse(initData string) (InitData, error) {
 	return d, nil
 }
 
-func GenerateSHA256(input string) string {
-	hash := sha256.New()
+func GenerateMD5(input string) string {
+	hash := md5.New()
 	hash.Write([]byte(input))
 	return hex.EncodeToString(hash.Sum(nil))
 }

src/components/Message.vue

@@ -3,22 +3,28 @@ import { ref } from 'vue'
 
 import axios from '@/plugins/axios'
 import { decryptE2EPacket, createE2EPacket } from '@/cryptography/DiffieHellman'
+import * as RSA from '@/cryptography/RSA'
+import * as AES from '@/cryptography/AES'
 
 import Time from '@/components/UI/Time.vue'
 import Button from '@/components/UI/Button.vue'
 import Textarea from '@/components/UI/Textarea.vue'
 
 const props = defineProps<{
-  id: number
-  text: string
-  time: string
-  owner: boolean
-  mark: boolean
-  canReplay: boolean
-  sender_public_key: string
-  quote?: {
+  message: {
     id: number
     content: string
+    time: string
+    owner: boolean
+    mark: boolean
+    can_replay: boolean
+    session_id: number
+    session_key?: string
+    sender_public_key: string
+    quote?: {
+      id: number
+      content: string
+    }
   }
 }>()
 
@@ -29,32 +35,34 @@ const replaySent = ref(false)
 const vDecrypt = {
   mounted: async (el: HTMLParagraphElement) => {
     try {
-      const isQuote = !!el.getAttribute('quote')
-      if (props.owner) {
-        window.Telegram.WebApp.CloudStorage.getItem(
-          'receive_private_key',
-          async (error, privateKey) => {
-            const decryptedMsg = await decryptE2EPacket(
-              privateKey!,
-              props.sender_public_key,
-              el.innerText,
-            )
-            el.innerText = decryptedMsg!
-          },
-        )
-      } else {
-        window.Telegram.WebApp.CloudStorage.getItem(
-          'send_private_key',
-          async (error, privateKey) => {
-            const decryptedMsg = await decryptE2EPacket(
-              privateKey!,
-              props.sender_public_key,
-              el.innerText,
-            )
-            el.innerText = decryptedMsg!
-          },
-        )
-      }
+      window.Telegram.WebApp.CloudStorage.getItem(
+        String(props.message.session_id),
+        async (error, sessionKey) => {
+          if (!sessionKey) {
+            window.Telegram.WebApp.CloudStorage.getItem("private_key", async (error, value) => {
+              const decryptedSessionKey = await RSA.decrypt(props.message.session_key!, value!)
+              window.Telegram.WebApp.CloudStorage.setItem(String(props.message.session_id), decryptedSessionKey)
+
+              try {
+                const decryptedMsg = await AES.decrypt(el.innerText, decryptedSessionKey!)
+                el.innerText = decryptedMsg!
+              } catch (error) {
+                alert(error)
+                el.innerText = 'خطا در رمزگشایی!'
+              }
+            })
+          } else {
+            try {
+              const decryptedMsg = await AES.decrypt(el.innerText, sessionKey!)
+              el.innerText = decryptedMsg!
+            } catch (error) {
+              alert(error)
+              el.innerText = 'خطا در رمزگشایی!'
+            }
+          }
+
+        },
+      )
     } catch (error) {
       alert(error)
       el.innerText = 'خطا در رمزگشایی!'
@@ -71,51 +79,36 @@ const vFocus = {
 function Submit() {
   if (!replayMessage.value) return
 
-  axios.get(`/get-key/${props.id}`).then(async ({ data: key }) => {
-    window.Telegram.WebApp.CloudStorage.getItem(
-      props.owner ? 'receive_private_key' : 'send_private_key',
-      async (error, privateKey) => {
-        const encryptedMsg = await createE2EPacket(
-          key,
-          privateKey!,
-          replayMessage.value,
-        )
-        axios
-          .post('/replay-message', {
-            message_id: props.id,
-            message: encryptedMsg,
-          })
-          .then(() => {
-            replaying.value = false
-            replayMessage.value = ''
-            replaySent.value = true
+  window.Telegram.WebApp.CloudStorage.getItem(String(props.message.session_id), async (error, sessionKey) => {
+    const encryptedMsg = await AES.encrypt(replayMessage.value, sessionKey!)
+    axios
+      .post('/replay-message', {
+        message_id: props.message.id,
+        message: encryptedMsg,
+      })
+      .then(() => {
+        replaying.value = false
+        replayMessage.value = ''
+        replaySent.value = true
 
-            setTimeout(() => (replaySent.value = false), 1500)
-          })
-      },
-    )
+        setTimeout(() => (replaySent.value = false), 1500)
+      }).catch((err: any) => {
+        alert(err)
+      })
   })
 }
 </script>
 <template>
-  <div
-    class="flex flex-col bg-[#ffffff] px-4 pt-3 pb-4 rounded-lg shadow-sm"
-    :class="mark && ['border-2 border-[#119af5]']"
-  >
-    <Time :value="time" class="text-gray-400 text-end text-sm"></Time>
-    <p
-      v-if="quote?.content"
-      class="border-r-4 rounded-md border-r-blue-500 pr-2 py-2 mt-2 truncate w-full"
-      style="background-color: rgba(137, 207, 240, 0.3)"
-      quote="true"
-      v-decrypt
-    >
-      {{ quote.content }}
+  <div class="flex flex-col bg-[#ffffff] px-4 pt-3 pb-4 rounded-lg shadow-sm">
+    <Time :value="message.time" class="text-gray-400 text-end text-sm"></Time>
+    <p v-if="message.quote?.content" class="border-r-4 rounded-md border-r-blue-500 pr-2 py-2 mt-2 truncate w-full"
+      style="background-color: rgba(137, 207, 240, 0.3)" quote="true" v-decrypt>
+      {{ message.quote.content }}
     </p>
 
-    <p class="break-words py-2" dir="auto" v-decrypt>{{ text }}</p>
+    <p class="break-words py-2" dir="auto" v-decrypt>{{ message.content }}</p>
 
-    <template v-if="canReplay">
+    <template v-if="message.can_replay">
       <div v-if="!replaying" class="flex justify-end text-gray-400 text-end">
         <div class="flex items-center cursor-pointer" @click="replaying = true">
           <span class="ml-1 text-sm">پاسخ</span>
@@ -124,16 +117,9 @@ function Submit() {
       </div>
 
       <div v-else class="flex flex-col mt-4">
-        <Textarea
-          v-model="replayMessage"
-          placeholder="پاسخ شما..."
-          v-focus
-        ></Textarea>
+        <Textarea v-model="replayMessage" placeholder="پاسخ شما..." v-focus></Textarea>
         <Button :block="true" class="mt-4" @click="Submit">ارسال</Button>
-        <p
-          class="text-center pt-4 text-[#119af5] font-bold cursor-pointer"
-          @click="replaying = false"
-        >
+        <p class="text-center pt-4 text-[#119af5] font-bold cursor-pointer" @click="replaying = false">
           بیخیال
         </p>
       </div>

src/components/Settings.vue

@@ -28,13 +28,13 @@ async function generateKeys() {
 
   await generateKeyPair().then(({ privateKey, publicKey }) => {
     receivePublicKey = publicKey
-    window.Telegram.WebApp.CloudStorage.setItem("receive_private_key", privateKey)
+    window.Telegram.WebApp.CloudStorage.setItem("private_key", privateKey)
 
   })
 
   await generateKeyPair().then(({ privateKey, publicKey }) => {
     sendPublicKey = publicKey
-    window.Telegram.WebApp.CloudStorage.setItem("send_private_key", privateKey)
+    window.Telegram.WebApp.CloudStorage.setItem("private_key", privateKey)
   })
 
   axios
@@ -43,11 +43,10 @@ async function generateKeys() {
       receive_public_key: receivePublicKey!,
     })
     .then(({ data }) => {
-      userStore.user.sendPublicKey = data.send_public_key
-      userStore.user.receivePublicKey = data.receive_public_key
+      userStore.user.publicKey = data.public_key
 
-      window.Telegram.WebApp.CloudStorage.setItem("send_public_key", data.send_public_key)
-      window.Telegram.WebApp.CloudStorage.setItem("receive_public_key", data.receive_public_key)
+      window.Telegram.WebApp.CloudStorage.setItem("public_key", data.send_public_key)
+      window.Telegram.WebApp.CloudStorage.setItem("public_key", data.receive_public_key)
     })
     .finally(() => {
       loading.value = false
@@ -68,8 +67,7 @@ function importHandler(event: Event) {
         receive_public_key: keys.receivePublicKey,
       })
       .then(({ data }) => {
-        userStore.user.sendPublicKey = data.send_public_key
-        userStore.user.receivePublicKey = data.receive_public_key
+        userStore.user.publicKey = data.public_key
 
         window.Telegram.WebApp.CloudStorage.setItem("send_public_key", data.send_public_key)
         window.Telegram.WebApp.CloudStorage.setItem("receive_public_key", data.receive_public_key)