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

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