You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

helpers.py 3.1KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. import base64
  2. import hashlib
  3. import hmac
  4. import os
  5. import time
  6. from configs import dify_config
  7. def get_signed_file_url(upload_file_id: str) -> str:
  8. url = f"{dify_config.FILES_URL}/files/{upload_file_id}/file-preview"
  9. timestamp = str(int(time.time()))
  10. nonce = os.urandom(16).hex()
  11. key = dify_config.SECRET_KEY.encode()
  12. msg = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
  13. sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
  14. encoded_sign = base64.urlsafe_b64encode(sign).decode()
  15. return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}"
  16. def get_signed_file_url_for_plugin(filename: str, mimetype: str, tenant_id: str, user_id: str) -> str:
  17. # Plugin access should use internal URL for Docker network communication
  18. base_url = dify_config.INTERNAL_FILES_URL or dify_config.FILES_URL
  19. url = f"{base_url}/files/upload/for-plugin"
  20. timestamp = str(int(time.time()))
  21. nonce = os.urandom(16).hex()
  22. key = dify_config.SECRET_KEY.encode()
  23. msg = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}"
  24. sign = hmac.new(key, msg.encode(), hashlib.sha256).digest()
  25. encoded_sign = base64.urlsafe_b64encode(sign).decode()
  26. return f"{url}?timestamp={timestamp}&nonce={nonce}&sign={encoded_sign}&user_id={user_id}&tenant_id={tenant_id}"
  27. def verify_plugin_file_signature(
  28. *, filename: str, mimetype: str, tenant_id: str, user_id: str, timestamp: str, nonce: str, sign: str
  29. ) -> bool:
  30. data_to_sign = f"upload|{filename}|{mimetype}|{tenant_id}|{user_id}|{timestamp}|{nonce}"
  31. secret_key = dify_config.SECRET_KEY.encode()
  32. recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  33. recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()
  34. # verify signature
  35. if sign != recalculated_encoded_sign:
  36. return False
  37. current_time = int(time.time())
  38. return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT
  39. def verify_image_signature(*, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
  40. data_to_sign = f"image-preview|{upload_file_id}|{timestamp}|{nonce}"
  41. secret_key = dify_config.SECRET_KEY.encode()
  42. recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  43. recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()
  44. # verify signature
  45. if sign != recalculated_encoded_sign:
  46. return False
  47. current_time = int(time.time())
  48. return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT
  49. def verify_file_signature(*, upload_file_id: str, timestamp: str, nonce: str, sign: str) -> bool:
  50. data_to_sign = f"file-preview|{upload_file_id}|{timestamp}|{nonce}"
  51. secret_key = dify_config.SECRET_KEY.encode()
  52. recalculated_sign = hmac.new(secret_key, data_to_sign.encode(), hashlib.sha256).digest()
  53. recalculated_encoded_sign = base64.urlsafe_b64encode(recalculated_sign).decode()
  54. # verify signature
  55. if sign != recalculated_encoded_sign:
  56. return False
  57. current_time = int(time.time())
  58. return current_time - int(timestamp) <= dify_config.FILES_ACCESS_TIMEOUT