Commit 4f79d572 authored by gurusrinadh.dasari's avatar gurusrinadh.dasari

Update main.py

parent 18af1ca4
...@@ -20,16 +20,17 @@ from jwt_util import Secrets ...@@ -20,16 +20,17 @@ from jwt_util import Secrets
ap = argparse.ArgumentParser() ap = argparse.ArgumentParser()
def create_token(user_id, ip, token, age=Secrets.LOCK_OUT_TIME_MINS): def create_token(user_id, ip, token, age=Secrets.LOCK_OUT_TIME_MINS, project_id="project_099"):
""" """
This method is to create a cookie This method is to create a cookie
""" """
try: try:
uid = str(uuid.uuid4()).replace("-", "") uid = str(uuid.uuid4()).replace("-", "")
payload = {"ip": ip, "user_id": user_id, "token": token, "uid": uid, "age": age} payload = {"ip": ip, "user_id": user_id, "token": token, "uid": uid, "age": age,
"project_id": project_id}
exp = datetime.now(timezone.utc) + timedelta(minutes=age) exp = datetime.utcnow() + timedelta(minutes=age)
_extras = {"iss": Secrets.issuer, "exp": exp} _extras = {"iss": Secrets.issuer, "exp": exp}
_payload = {**payload, **_extras} _payload = {**payload, **_extras}
...@@ -81,6 +82,13 @@ if __name__ == "__main__": ...@@ -81,6 +82,13 @@ if __name__ == "__main__":
default='127.0.0.1', default='127.0.0.1',
help="Ip details for the login-token", help="Ip details for the login-token",
) )
ap.add_argument(
"--PROJECT_ID",
"-pid",
required=False,
default='project_099',
help="Project ID for the token generation..",
)
arguments = vars(ap.parse_args()) arguments = vars(ap.parse_args())
print(f"{arguments['REDIS_URI']}:{arguments['TOKEN_AGE']}:{arguments['REDIS_DB']}") print(f"{arguments['REDIS_URI']}:{arguments['TOKEN_AGE']}:{arguments['REDIS_DB']}")
...@@ -88,11 +96,12 @@ if __name__ == "__main__": ...@@ -88,11 +96,12 @@ if __name__ == "__main__":
TOKEN_AGE = arguments['TOKEN_AGE'] TOKEN_AGE = arguments['TOKEN_AGE']
REDIS_DB = arguments['REDIS_DB'] REDIS_DB = arguments['REDIS_DB']
USER_ID = arguments['USER_ID'] USER_ID = arguments['USER_ID']
PROJECT_ID = arguments['PROJECT_ID']
ip_details = arguments['HOST'] ip_details = arguments['HOST']
if not all([REDIS_URI, TOKEN_AGE, REDIS_DB]): if not all([REDIS_URI, TOKEN_AGE, REDIS_DB]):
sys.exit() sys.exit()
login_db = redis.from_url(REDIS_URI, db=int(REDIS_DB), decode_responses=True) login_db = redis.from_url(REDIS_URI, db=int(REDIS_DB), decode_responses=True)
age_in_minutes = int(TOKEN_AGE) * 24 * 60 age_in_minutes = int(TOKEN_AGE) * 24 * 60
token_details = create_token(user_id=USER_ID, ip=ip_details, token=Secrets.token, age=age_in_minutes) token_details = create_token(user_id=USER_ID, ip=ip_details, token=Secrets.token, age=age_in_minutes, project_id=PROJECT_ID)
typer.echo(typer.style(f"Token Generated for {TOKEN_AGE} days is {token_details}", fg=typer.colors.GREEN)) typer.echo(typer.style(f"Token Generated for {TOKEN_AGE} days is {token_details}", fg=typer.colors.GREEN))
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment