utils.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. #-*- coding:utf-8 -*-
  2. import json
  3. import requests
  4. from flask import g, redirect, session, abort, request
  5. from functools import wraps
  6. from rrd import config
  7. from rrd import corelib
  8. from rrd.utils import randbytes
  9. from rrd.model.user import User, UserToken
  10. from rrd.utils.logger import logging
  11. log = logging.getLogger(__file__)
  12. def remote_ip():
  13. if not request.headers.getlist("X-Forward-For"):
  14. return request.remote_addr
  15. else:
  16. return request.headers.getlist("X-Forward-For")[0]
  17. def require_login(redir="/auth/login"):
  18. def _(f):
  19. @wraps(f)
  20. def __(*a, **kw):
  21. if not g.user:
  22. return redirect(redir or "/auth/login")
  23. return f(*a, **kw)
  24. return __
  25. return _
  26. def require_login_abort(status_code=403, msg="login first"):
  27. def _(f):
  28. @wraps(f)
  29. def __(*a, **kw):
  30. if not g.user:
  31. return abort(status_code, msg)
  32. return f(*a, **kw)
  33. return __
  34. return _
  35. def require_login_json(json_msg={"ok":False, "msg":"login first"}):
  36. def _(f):
  37. @wraps(f)
  38. def __(*a, **kw):
  39. if not g.user:
  40. return json.dumps(json_msg)
  41. return f(*a, **kw)
  42. return __
  43. return _
  44. def set_user_cookie(user_token, session_):
  45. if not user_token:
  46. return None
  47. session_[config.SITE_COOKIE] = "%s:%s" % (user_token.name, user_token.sig)
  48. def clear_user_cookie(session_):
  49. session_[config.SITE_COOKIE] = ""
  50. def get_usertoken_from_session(session_):
  51. if config.SITE_COOKIE in session_:
  52. cookies = session_[config.SITE_COOKIE]
  53. if not cookies:
  54. return None
  55. name, sig = cookies.split(":")
  56. return UserToken(name, sig)
  57. def get_current_user_profile(user_token):
  58. if not user_token:
  59. return
  60. h = {"Content-type": "application/json"}
  61. r = corelib.auth_requests("GET", "%s/user/current" %config.API_ADDR, headers=h)
  62. if r.status_code != 200:
  63. return
  64. j = r.json()
  65. return User(j["id"], j["name"], j["cnname"], j["email"], j["phone"], j["im"], j["qq"], j["role"])
  66. def logout_user(user_token):
  67. if not user_token:
  68. return
  69. r = corelib.auth_requests("GET", "%s/user/logout" %config.API_ADDR)
  70. if r.status_code != 200:
  71. raise Exception("%s:%s" %(r.status_code, r.text))
  72. clear_user_cookie(session)
  73. def login_user(name, password):
  74. params = {
  75. "name": name,
  76. "password": password,
  77. }
  78. r = requests.post("%s/user/login" %config.API_ADDR, data=params)
  79. if r.status_code != 200:
  80. raise Exception("{} : {}".format(r.status_code, r.text))
  81. j = r.json()
  82. ut = UserToken(j["name"], j["sig"])
  83. set_user_cookie(ut, session)
  84. return ut
  85. def ldap_login_user(name, password):
  86. import ldap
  87. if not config.LDAP_ENABLED:
  88. raise Exception("ldap not enabled")
  89. bind_dn = config.LDAP_BINDDN_FMT
  90. try:
  91. bind_dn = config.LDAP_BINDDN_FMT %name
  92. except TypeError: pass
  93. search_filter = config.LDAP_SEARCH_FMT
  94. try:
  95. search_filter = config.LDAP_SEARCH_FMT %name
  96. except TypeError: pass
  97. cli = None
  98. try:
  99. ldap_server = config.LDAP_SERVER if config.LDAP_SERVER.startswith("ldap://") else "ldap://%s" %config.LDAP_SERVER
  100. log.debug("bind_dn=%s base_dn=%s filter=%s attrs=%s" %(bind_dn, config.LDAP_BASE_DN, search_filter, config.LDAP_ATTRS))
  101. cli = ldap.initialize(ldap_server)
  102. cli.bind_s(bind_dn, password)
  103. result = cli.search_s(config.LDAP_BASE_DN, ldap.SCOPE_SUBTREE, search_filter, config.LDAP_ATTRS)
  104. log.debug("ldap result: %s" % result)
  105. d = result[0][1]
  106. email = d['mail'][0]
  107. cnname = d['cn'][0]
  108. if 'telephoneNumber' in d:
  109. phone = d['telephoneNumber'] and d['telephoneNumber'][0] or ""
  110. else:
  111. phone = ""
  112. return {
  113. "name": name,
  114. "password": password,
  115. "cnname": cnname,
  116. "email": email,
  117. "phone": phone,
  118. }
  119. except ldap.LDAPError as e:
  120. cli and cli.unbind_s()
  121. raise e
  122. except (IndexError, KeyError) as e:
  123. raise e
  124. finally:
  125. cli and cli.unbind_s()