Spaces:
Running
Running
File size: 5,727 Bytes
3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 75bcdb3 3c37508 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
#!/usr/bin/env python3
"""
Switch Trackio Space HF_TOKEN from Write Token to Read Token
This script switches the HF_TOKEN secret in a Trackio Space from a write token
to a read token after the experiment is complete, for security purposes.
The space uses only HF_TOKEN, which starts as write token and gets switched to read token.
"""
import os
import sys
import json
from typing import Optional, Tuple
from huggingface_hub import HfApi
def validate_token_permissions(token: str) -> Tuple[bool, str, Optional[str]]:
"""
Validate token and determine its permission level.
Args:
token (str): The Hugging Face token to validate
Returns:
Tuple[bool, str, Optional[str]]:
- success: True if token is valid
- permission_level: "read" or "write"
- username: The username associated with the token
"""
try:
api = HfApi(token=token)
user_info = api.whoami()
# Extract username
username = user_info.get("name", user_info.get("username"))
# Test write permissions by trying to access a test repository
# We'll use a simple test - try to get repo info for a public repo
try:
# Try to access a public dataset to test read permissions
api.dataset_info("huggingface-course/documentation-tutorial")
# For write permissions, we'll assume the token has write access
# since we can't easily test write permissions without creating something
# In practice, write tokens are typically provided by users who know
# they have write access
return True, "write", username
except Exception as e:
# If we can't access even a public dataset, it's likely a read token
return True, "read", username
except Exception as e:
error_msg = str(e)
if "401" in error_msg or "unauthorized" in error_msg.lower():
return False, "invalid", None
else:
return False, "error", None
def switch_space_token(space_id: str, read_token: str, write_token: str) -> bool:
"""
Switch the HF_TOKEN secret in a Trackio Space from write to read token.
Args:
space_id (str): The space ID (username/space-name)
read_token (str): The read token to set as new HF_TOKEN
write_token (str): The write token (for authentication to update the space)
Returns:
bool: True if successful, False otherwise
"""
try:
# Validate both tokens
print("🔍 Validating tokens...")
write_valid, write_perm, write_user = validate_token_permissions(write_token)
read_valid, read_perm, read_user = validate_token_permissions(read_token)
if not write_valid:
print(f"❌ Write token validation failed")
return False
if not read_valid:
print(f"❌ Read token validation failed")
return False
if write_user != read_user:
print(f"❌ Token mismatch: write token user ({write_user}) != read token user ({read_user})")
return False
print(f"✅ Tokens validated successfully")
print(f" Write token: {write_perm} permissions for {write_user}")
print(f" Read token: {read_perm} permissions for {read_user}")
# Use the write token to update the space (since we need write access)
api = HfApi(token=write_token)
# Update the HF_TOKEN secret in the space from write token to read token
try:
api.add_space_secret(
repo_id=space_id,
key="HF_TOKEN",
value=read_token,
description="Hugging Face token for dataset access (switched from write to read for security)"
)
print(f"✅ Successfully switched HF_TOKEN from write to read token in space: {space_id}")
print(f"🔒 Space now uses read-only permissions for enhanced security")
return True
except Exception as e:
print(f"❌ Failed to update HF_TOKEN secret: {e}")
return False
except Exception as e:
print(f"❌ Error switching HF_TOKEN: {e}")
return False
def main():
"""Main function to switch tokens."""
print("🔄 Trackio Space Token Switch")
print("=" * 40)
# Get arguments
if len(sys.argv) >= 4:
space_id = sys.argv[1]
read_token = sys.argv[2]
write_token = sys.argv[3]
else:
print("Usage: python switch_to_read_token.py <space_id> <read_token> <write_token>")
print("Example: python switch_to_read_token.py username/trackio-monitoring read_token write_token")
sys.exit(1)
# Validate space_id format
if "/" not in space_id:
print("❌ Invalid space_id format. Use: username/space-name")
sys.exit(1)
# Switch tokens
success = switch_space_token(space_id, read_token, write_token)
if success:
print("\n✅ HF_TOKEN switch completed successfully!")
print(f"📊 Space: {space_id}")
print("🔒 HF_TOKEN now uses read-only permissions for enhanced security")
print("💡 The space can still read datasets but cannot write to repositories")
print("🎯 Training is complete - space is now secure for monitoring")
else:
print("\n❌ HF_TOKEN switch failed!")
print("Please check your tokens and try again.")
sys.exit(1)
if __name__ == "__main__":
main() |