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.3KB

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