Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 76ffd64

Browse files
nodeweenodewee
committed
Add encryption support for sending messages
1 parent 62a446e commit 76ffd64

File tree

6 files changed

+190
-155
lines changed

6 files changed

+190
-155
lines changed

‎mixinsdk/api/message.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import time
12
from io import FileIO
23
from typing import List, Union
34

@@ -16,9 +17,11 @@ def send_messages(self, messages: Union[List[dict], dict]):
1617
A maximum of 100 messages can be sent in batch each time,
1718
and the message body cannot exceed 128Kb.
1819
"""
19-
2020
return self._http.post("/messages", messages)
2121

22+
def send_encrypted_messages(self, messages: list):
23+
return self._http.post("/encrypted_messages", messages)
24+
2225
def create_attachment(self) -> dict:
2326
"""After creating action, then upload the attachment to upload_url,
2427
and then the attachment_id can be used sending images,

‎mixinsdk/clients/_message.py‎

Lines changed: 70 additions & 107 deletions
Original file line numberDiff line numberDiff line change
@@ -1,50 +1,54 @@
11
import base64
22
import json
3-
import os
3+
import secrets
44
import uuid
55
from typing import List, Union
66

77
import nacl.bindings
8-
from cryptography.hazmat.primitives.asymmetric import ed25519
8+
import nacl.signing
9+
from cryptography.hazmat.backends import default_backend
10+
from cryptography.hazmat.primitives import serialization
11+
from cryptography.hazmat.primitives.asymmetric import ed25519, x25519
912
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
13+
from nacl.encoding import RawEncoder
14+
from nacl.signing import SigningKey
1015

11-
from mixinsdk.types.user import UserSession
1216
from mixinsdk.utils import base64_pad_equal_sign
1317

14-
from ..types.message import MESSAGE_CATEGORIES
18+
from ..utils import base64_pad_equal_sign
1519

1620

1721
def parse_message_data(
18-
data: str, category: str, app_session_id: str, app_private_key: bytes
22+
data_b64_str: str, category: str, app_session_id: str, app_private_key: bytes
1923
) -> Union[dict, str]:
2024
"""
21-
- parse message data (Base64 encoded string) to str or dict
22-
- decrypt message data if category is ENCRYPTED_
25+
- parse message data to str or dict. if category is ENCRYPTED_*, will decrypt message data first.
2326
2427
Returns: data_parsed
2528
"""
26-
if not data:
29+
if not data_b64_str:
2730
return ""
2831

2932
if category.startswith("ENCRYPTED_"):
30-
d = decrypt_message_data(data, app_session_id, app_private_key)
33+
d = decrypt_message_data(data_b64_str, app_session_id, app_private_key)
3134
# print("\ndecrypted:", d)
3235
else:
33-
d = base64.b64decode(data).decode()
36+
d = base64.b64decode(data_b64_str).decode()
3437

3538
if category.endswith(("_TEXT", "_POST")):
3639
return d
3740
try:
3841
return json.loads(d)
3942
except json.JSONDecodeError:
40-
print(f"Failed to json decode data: {d}")
43+
print(f"Failed to json decode data_b64_str: {d}")
4144

4245

4346
def decrypt_message_data(data_b64_str: str, app_session_id: str, private: bytes):
4447
data_bytes = base64.b64decode(base64_pad_equal_sign(data_b64_str))
45-
size = 16 + 48 # session id bytes + encrypted key bytes size
48+
size = 16 + 48 # session id bytes + shared key bytes size
4649
total = len(data_bytes)
47-
if total < 1 + 2 + 32 + size + 12:
50+
# bytes([1]) + session_len + pub_key + session_id + shared_key + nonce + ...data
51+
if total < 1 + 2 + 32 + 16 + 48 + 12:
4852
raise ValueError("Invalid message data")
4953
session_length = int.from_bytes(data_bytes[1:3], byteorder="little")
5054
prefixSize = 35 + session_length * size
@@ -63,113 +67,72 @@ def decrypt_message_data(data_b64_str: str, app_session_id: str, private: bytes)
6367
block_size = 16
6468
iv = data_bytes[i + 16 : i + 16 + block_size]
6569
key = data_bytes[i + 16 + block_size : i + size]
66-
cipher = Cipher(algorithms.AES(dst), modes.CBC(iv))
67-
encryptor = cipher.decryptor()
68-
key = encryptor.update(key) # do not: + encryptor.finalize()
70+
decryptor = Cipher(algorithms.AES(dst), modes.CBC(iv)).decryptor()
71+
key = decryptor.update(key) # do not: + decryptor.finalize()
6972
key = key[:16]
7073
break
7174

7275
if len(key) != 16:
7376
raise ValueError("Invalid key")
7477

75-
encrypted_data = data_bytes[prefixSize + 12 : -16] # ? -16 是根据结果错误输出来的，go 的代码没有-16
78+
encrypted_data = data_bytes[prefixSize + 12 : -16] # ?len(finalize() + .tag) = 16
7679

7780
nonce = data_bytes[prefixSize : prefixSize + 12]
78-
cipher = Cipher(algorithms.AES(key), modes.GCM(nonce))
79-
encryptor = cipher.decryptor()
80-
plain_text = encryptor.update(encrypted_data)
81+
decryptor = Cipher(algorithms.AES(key), modes.GCM(nonce)).decryptor()
82+
plain_text = decryptor.update(encrypted_data)
8183
return plain_text.decode()
8284

8385

8486
def encrypt_message_data(
85-
data_b64_str: str, user_sessions: List[UserSession], app_private_key: bytes
87+
data_bytes: bytes, user_sessions: List[dict], app_private_key: bytes
8688
):
87-
data_bytes = base64.b64decode(data_b64_str)
88-
8989
"""
90-
key := make([]byte, 16)
91-
_, err = rand.Read(key)
92-
if err != nil {
93-
return "", err
94-
}
95-
nonce := make([]byte, 12)
96-
_, err = rand.Read(nonce)
97-
if err != nil {
98-
return "", err
99-
}
100-
block, err := aes.NewCipher(key)
101-
if err != nil {
102-
return "", err
103-
}
104-
aesgcm, err := cipher.NewGCM(block)
105-
if err != nil {
106-
return "", err
107-
}
108-
ciphertext := aesgcm.Seal(nil, nonce, data_bytes, nil)
90+
session struct: {user_id:uuid str, session_id:uuid str, public_key:str}
10991
"""
110-
block_size = 16
111-
key = os.urandom(block_size)
112-
nonce = os.urandom(12)
113-
cipher = Cipher(algorithms.AES(key), modes.GCM(nonce))
114-
encryptor = cipher.encryptor()
115-
ciphertext = encryptor.update(data_bytes) # do not: + encryptor.finalize()
116-
117-
session_length = len(user_sessions)
118-
session_length = session_length.to_bytes(2, byteorder="little")
119-
120-
# private := ed25519.PrivateKey(privateBytes)
121-
# pub, _ := PublicKeyToCurve25519(ed25519.PublicKey(private[32:]))
122-
private = ed25519.Ed25519PrivateKey().from_private_bytes(app_private_key)
123-
public = private.public_key()
124-
priv_curve25519 = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(private)
125-
pub_curve25519 = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(public)
126-
127-
session_bytes = bytearray()
128-
# for _, s := range sessions {
92+
93+
key = secrets.token_bytes(16)
94+
nonce = secrets.token_bytes(12)
95+
96+
encryptor = Cipher(algorithms.AES(key), modes.GCM(nonce)).encryptor()
97+
main_ciphertext = encryptor.update(data_bytes)
98+
main_ciphertext += encryptor.finalize() + encryptor.tag
99+
100+
# ed25519 private key -> cureve25519 public key
101+
_pk = nacl.bindings.crypto_sign_ed25519_sk_to_pk(app_private_key)
102+
curve25519_pubkey = nacl.bindings.crypto_sign_ed25519_pk_to_curve25519(_pk)
103+
104+
# ed25519 private key -> curve25519 private key
105+
curve25519_privkey = nacl.bindings.crypto_sign_ed25519_sk_to_curve25519(
106+
app_private_key
107+
)
108+
109+
session_len = len(user_sessions).to_bytes(2, byteorder="little")
110+
sessions_bytes = b""
111+
112+
padding = 16 - len(key) % 16
113+
pad_text = bytes([padding] * padding)
114+
shared = key + pad_text
115+
129116
for s in user_sessions:
130-
# clientPublic, err := base64.RawURLEncoding.DecodeString(s.PublicKey)
131-
client_public = base64.b64decode(s.public_key)
132-
# var dst, priv, clientPub [32]byte
133-
# copy(clientPub[:], clientPublic[:])
134-
client_pub = client_public[:32]
135-
136-
# curve25519.ScalarMult(&dst, &priv, &clientPub)
137-
dst = nacl.bindings.crypto_scalarmult(priv_curve25519, client_pub)
138-
139-
# padding := aes.BlockSize - len(key)%aes.BlockSize
140-
padding = block_size - len(key) % block_size
141-
padtext = bytes(padding) * padding
142-
# copy(shared[:], key[:])
143-
shared = key[:]
144-
# shared = append(shared, padtext...)
145-
shared += padtext
146-
# ciphertext := make([]byte, aes.BlockSize+len(shared))
147-
# iv := ciphertext[:aes.BlockSize]
148-
iv = os.urandom(block_size + len(shared))
149-
# mode := cipher.NewCBCEncrypter(block, iv)
150-
# mode.CryptBlocks(ciphertext[aes.BlockSize:], shared)
151-
# block, err := aes.NewCipher(dst[:])
152-
cipher = Cipher(algorithms.AES(dst), modes.CBC(iv))
153-
encryptor = cipher.encryptor()
154-
ciphertext = encryptor.update(shared)
155-
# id, err := UuidFromString(s.SessionID)
156-
id = uuid.UUID(s.session_id)
157-
# sessionsBytes = append(sessionsBytes, id.Bytes()...)
158-
# sessionsBytes = append(sessionsBytes, ciphertext...)
159-
session_bytes.append(id.bytes)
160-
session_bytes.append(ciphertext)
161-
162-
# result := []byte{1}
163-
result = bytearray()
164-
# result = append(result, sessionLen[:]...)
165-
result.append(session_length)
166-
# result = append(result, pub[:]...)
167-
result.append(pub_curve25519)
168-
# result = append(result, sessionsBytes...)
169-
result.extend(session_bytes)
170-
# result = append(result, nonce[:]...)
171-
result.extend(nonce)
172-
# result = append(result, ciphertext...)
173-
result.extend(ciphertext)
174-
# return base64.RawURLEncoding.EncodeToString(result), nil
175-
return base64.b64encode(result)
117+
client_pub = base64.urlsafe_b64decode(base64_pad_equal_sign(s["public_key"]))
118+
dst = nacl.bindings.crypto_scalarmult(curve25519_privkey, client_pub)
119+
120+
iv = secrets.token_bytes(16)
121+
encryptor = Cipher(algorithms.AES(dst), modes.CBC(iv)).encryptor()
122+
ciphertext = iv + encryptor.update(shared) # do not: + encryptor.finalize()
123+
124+
id = uuid.UUID(s["session_id"]).bytes
125+
sessions_bytes += id + ciphertext
126+
127+
result = (
128+
bytes([1])
129+
+ session_len
130+
+ curve25519_pubkey
131+
+ sessions_bytes
132+
+ nonce
133+
+ main_ciphertext
134+
)
135+
encoded = base64.urlsafe_b64encode(result).decode("utf-8")
136+
encoded = encoded.rstrip("=") # remove padding
137+
138+
return encoded

‎mixinsdk/clients/client_blaze.py‎

Lines changed: 19 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import base64
23
import gzip
34
import json
45
import logging
@@ -11,14 +12,12 @@
1112

1213
import websockets
1314
import websockets.client
14-
from cryptography.hazmat.primitives.asymmetric import ed25519
1515

16-
from mixinsdk.clients._message import encrypt_message_data
17-
from mixinsdk.types.user import UserProfile
16+
from mixinsdk.types.user import UserProfile, UserSession
1817

1918
from ..constants import API_BASE_URLS
2019
from ..utils import get_conversation_id_of_two_users
21-
from ._message import parse_message_data
20+
from . import _message
2221
from ._sign import sign_authentication_token
2322
from .config import AppConfig
2423

@@ -86,9 +85,6 @@ def send_message(self, message: dict):
8685
- message, use types.message.pack_message() to make it
8786
"""
8887

89-
# TODO : depends on switch of encryption
90-
# message["data"]=
91-
9288
msg = {
9389
"id": str(uuid.uuid4()),
9490
"action": "CREATE_MESSAGE",
@@ -211,24 +207,25 @@ def _handle_message_done(future: asyncio.Future):
211207
break # exit the while loop
212208

213209
def parse_message_data(self, data: str, category: str):
214-
return parse_message_data(
210+
return _message.parse_message_data(
215211
data, category, self.config.session_id, self.config.private_key
216212
)
217213

218-
def encrypt_message_data(self, data_b64_str: str):
219-
pass
220-
# # data = base64.b64encode(b"hello world").decode("utf-8")
221-
# private = ed25519.Ed25519PrivateKey().from_private_bytes(
222-
# self.config..private_key
223-
# )
224-
# public = private.public_key()
225-
# user_session = UserSession(
226-
# self.config..client_id, self.config..session_id, public
227-
# )
228-
# data_encrypted = encrypt_message_data(
229-
# data, [user_session], self.config..private_key
230-
# )
231-
# data_b64_str = base64.b64encode(data_encrypted).decode("utf-8")
214+
def encrypt_message_data(self, b64encoded_data: str, conversation_id: str):
215+
data_bytes = base64.b64decode(b64encoded_data)
216+
user_sessions = self.get_conversation_user_sessions(conversation_id)
217+
# print("\n\nGot user_sessions:", user_sessions)
218+
encrypted_data = _message.encrypt_message_data(
219+
data_bytes, user_sessions, self.config.private_key
220+
)
221+
recipient_sessions = [
222+
{"session_id": session["session_id"]} for session in user_sessions
223+
]
224+
# print(recipient_sessions)
225+
# exit()
226+
checksum = self.generate_session_checksum(recipient_sessions)
227+
228+
return encrypted_data, recipient_sessions, checksum
232229

233230
def start_to_list_pending_message(self):
234231
if not self.ws:

0 commit comments

Comments
 (0)