Skip to content

Commit

Permalink
lib: add Quota model and the max. number of queries inside the CellBa…
Browse files Browse the repository at this point in the history
…se token, #TASK-4791, #TASK-4641
  • Loading branch information
jtarraga committed Jul 18, 2023
1 parent 24f7f32 commit 8dbb33f
Show file tree
Hide file tree
Showing 6 changed files with 165 additions and 42 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
import org.apache.commons.lang3.StringUtils;
import org.opencb.cellbase.app.cli.CommandExecutor;
import org.opencb.cellbase.app.cli.admin.AdminCliOptionsParser;
import org.opencb.cellbase.core.token.DataAccessTokenSources;
import org.opencb.cellbase.core.token.DataAccessToken;
import org.opencb.cellbase.core.token.DataAccessTokenManager;

import javax.crypto.spec.SecretKeySpec;
Expand Down Expand Up @@ -51,8 +51,8 @@ public void execute() {
try {
if (StringUtils.isNotEmpty(dataTokenCommandOptions.createWithDataSources)) {
// Create data token
DataAccessTokenSources dataSources = null;
dataSources = DataAccessTokenSources.parse(dataTokenCommandOptions.createWithDataSources);
DataAccessToken dataSources = null;
dataSources = DataAccessToken.parse(dataTokenCommandOptions.createWithDataSources);
String token = datManager.encode(dataTokenCommandOptions.organization, dataSources);
System.out.println("Data access token generated:\n" + token);
} else if (StringUtils.isNotEmpty(dataTokenCommandOptions.tokenToView)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,17 +24,27 @@
import java.util.HashMap;
import java.util.Map;

public class DataAccessTokenSources {
public class DataAccessToken {

private String version;
private Map<String, Long> sources;
private Long maxNumQueries;

public DataAccessTokenSources() {
public static final String CURRENT_VERSION = "1.0";
public static final Long MAX_NUM_ANOYMOUS_QUERIES = 1000000L;

public DataAccessToken() {
this(CURRENT_VERSION, new HashMap<>(), 0L);
}

public DataAccessToken(String version, Map<String, Long> sources) {
this(version, sources, 0L);
}

public DataAccessTokenSources(String version, Map<String, Long> sources) {
public DataAccessToken(String version, Map<String, Long> sources, Long maxNumQueries) {
this.version = version;
this.sources = sources;
this.maxNumQueries = maxNumQueries;
}

private static DateFormat dateFormatter = new SimpleDateFormat("dd/MM/yyyy");
Expand All @@ -43,32 +53,38 @@ public static DateFormat dateFormatter() {
return dateFormatter;
}

public static DataAccessTokenSources parse(String input) throws ParseException {
DataAccessTokenSources dataSources = new DataAccessTokenSources();
Map<String, Long> sources = new HashMap<>();
String[] split = input.split(",");
public static DataAccessToken parse(String sources) throws ParseException {
return parse(sources, null);
}

public static DataAccessToken parse(String sources, Long maxNumQueries) throws ParseException {
DataAccessToken dataAccessToken = new DataAccessToken();
Map<String, Long> sourcesMap = new HashMap<>();
String[] split = sources.split(",");
for (String source : split) {
String[] splits = source.split(":");
if (splits.length == 1) {
sources.put(splits[0], Long.MAX_VALUE);
sourcesMap.put(splits[0], Long.MAX_VALUE);
} else {
sources.put(splits[0], dateFormatter.parse(splits[1]).getTime());
sourcesMap.put(splits[0], dateFormatter.parse(splits[1]).getTime());
}
}

dataSources.setVersion("1.0");
if (MapUtils.isNotEmpty(sources)) {
dataSources.setSources(sources);
dataAccessToken.setVersion(CURRENT_VERSION);
if (MapUtils.isNotEmpty(sourcesMap)) {
dataAccessToken.setSources(sourcesMap);
}
dataAccessToken.setMaxNumQueries(maxNumQueries);

return dataSources;
return dataAccessToken;
}

@Override
public String toString() {
final StringBuilder sb = new StringBuilder("DataAccessTokenSources{");
final StringBuilder sb = new StringBuilder("DataAccessToken{");
sb.append("version='").append(version).append('\'');
sb.append(", sources=").append(sources);
sb.append(", maxNumQueries=").append(maxNumQueries);
sb.append('}');
return sb.toString();
}
Expand All @@ -77,7 +93,7 @@ public String getVersion() {
return version;
}

public DataAccessTokenSources setVersion(String version) {
public DataAccessToken setVersion(String version) {
this.version = version;
return this;
}
Expand All @@ -86,8 +102,17 @@ public Map<String, Long> getSources() {
return sources;
}

public DataAccessTokenSources setSources(Map<String, Long> sources) {
public DataAccessToken setSources(Map<String, Long> sources) {
this.sources = sources;
return this;
}

public Long getMaxNumQueries() {
return maxNumQueries;
}

public DataAccessToken setMaxNumQueries(Long maxNumQueries) {
this.maxNumQueries = maxNumQueries;
return this;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,22 +28,27 @@
import java.security.Key;
import java.util.*;

import static org.opencb.cellbase.core.token.DataAccessTokenSources.dateFormatter;
import static org.opencb.cellbase.core.token.DataAccessToken.MAX_NUM_ANOYMOUS_QUERIES;
import static org.opencb.cellbase.core.token.DataAccessToken.dateFormatter;

public class DataAccessTokenManager {
private SignatureAlgorithm algorithm;
private Key privateKey;
private Key publicKey;
private JwtParser jwtParser;

private String defaultToken;

private final Logger logger = LoggerFactory.getLogger(DataAccessTokenManager.class);

public static final int SECRET_KEY_MIN_LENGTH = 50;
public static final String VERSION_FIELD_NAME = "version";
public static final String SOURCES_FIELD_NAME = "sources";
public static final String MAX_NUM_QUERIES_FIELD_NAME = "maxNumQueries";

public DataAccessTokenManager(String key) {
this(SignatureAlgorithm.HS256.getValue(), new SecretKeySpec(TextCodec.BASE64.decode(key), SignatureAlgorithm.HS256.getJcaName()));
defaultToken = encode("ANONYMOUS", new DataAccessToken(DataAccessToken.CURRENT_VERSION, new HashMap<>(), MAX_NUM_ANOYMOUS_QUERIES));
}

public DataAccessTokenManager(String algorithm, Key secretKey) {
Expand All @@ -57,14 +62,15 @@ public DataAccessTokenManager() {
jwtParser = Jwts.parserBuilder().build();
}

public String encode(String organization, DataAccessTokenSources dat) {
public String encode(String organization, DataAccessToken dat) {
JwtBuilder jwtBuilder = Jwts.builder();

Map<String, Object> claims = new HashMap<>();
claims.put(VERSION_FIELD_NAME, dat.getVersion());
if (MapUtils.isNotEmpty(dat.getSources())) {
claims.put(SOURCES_FIELD_NAME, dat.getSources());
}
claims.put(MAX_NUM_QUERIES_FIELD_NAME, dat.getMaxNumQueries());

jwtBuilder.setClaims(claims)
.setSubject(organization)
Expand All @@ -74,8 +80,8 @@ public String encode(String organization, DataAccessTokenSources dat) {
return jwtBuilder.compact();
}

public DataAccessTokenSources decode(String token) {
DataAccessTokenSources dat = new DataAccessTokenSources();
public DataAccessToken decode(String token) {
DataAccessToken dat = new DataAccessToken();

Claims body = parse(token);
for (Map.Entry<String, Object> entry : body.entrySet()) {
Expand All @@ -87,6 +93,9 @@ public DataAccessTokenSources decode(String token) {
case SOURCES_FIELD_NAME:
dat.setSources((Map<String, Long>) body.get(key));
break;
case MAX_NUM_QUERIES_FIELD_NAME:
dat.setMaxNumQueries(((Integer)body.get(key)).longValue());
break;
default:
break;
}
Expand All @@ -96,26 +105,26 @@ public DataAccessTokenSources decode(String token) {
}

public String recode(String token) {
DataAccessTokenSources dataAccessTokenSources = decode(token);
if (MapUtils.isNotEmpty(dataAccessTokenSources.getSources())) {
DataAccessToken dataAccessToken = decode(token);
if (MapUtils.isNotEmpty(dataAccessToken.getSources())) {
Map<String, Long> sources = new HashMap<>();
for (Map.Entry<String, Long> entry : dataAccessTokenSources.getSources().entrySet()) {
for (Map.Entry<String, Long> entry : dataAccessToken.getSources().entrySet()) {
if (new Date().getTime() <= entry.getValue()) {
sources.put(entry.getKey(), entry.getValue());
}
}
dataAccessTokenSources.setSources(sources);
dataAccessToken.setSources(sources);
}

return encode(getOrganization(token), dataAccessTokenSources);
return encode(getOrganization(token), dataAccessToken);
}

public void validate(String token) {
parse(token);
}

public boolean hasExpiredSource(String source, String token) throws IllegalArgumentException {
DataAccessTokenSources dat = decode(token);
DataAccessToken dat = decode(token);
if (MapUtils.isNotEmpty(dat.getSources()) && dat.getSources().containsKey(source)) {
return (new Date().getTime() > dat.getSources().get(source));
}
Expand All @@ -133,7 +142,7 @@ public Set<String> getValidSources(String token, Set<String> init) throws Illega
}

if (StringUtils.isNotEmpty(token)) {
DataAccessTokenSources dat = decode(token);
DataAccessToken dat = decode(token);
if (MapUtils.isNotEmpty(dat.getSources())) {
for (Map.Entry<String, Long> entry : dat.getSources().entrySet()) {
if (new Date().getTime() <= entry.getValue()) {
Expand All @@ -160,6 +169,10 @@ public String getCreationDate(String token) {
return dateFormatter().format(parse.getIssuedAt());
}

public String getDefaultToken() {
return defaultToken;
}

public void display(String token) {
Claims body = parse(token);

Expand All @@ -173,6 +186,7 @@ public void display(String token) {
for (Map.Entry<String, Long> entry : sources.entrySet()) {
sb.append("\t- '").append(entry.getKey()).append("' until ").append(dateFormatter().format(entry.getValue())).append("\n");
}
sb.append("Max. num. queries: ").append(body.get(MAX_NUM_QUERIES_FIELD_NAME)).append("\n");

System.out.println(sb);
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,74 @@
/*
* Copyright 2015-2020 OpenCB
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/

package org.opencb.cellbase.core.token;

public class Quota {
private String token;
private String date; // date consists of year + month, e.g.: 202304
private long numQueries;

public Quota() {
}

public Quota(String token, String date, long numQueries) {
this.token = token;
this.date = date;
this.numQueries = numQueries;
}

@Override
public String toString() {
final StringBuilder sb = new StringBuilder("Quota{");
sb.append("token='").append(token).append('\'');
sb.append(", date='").append(date).append('\'');
sb.append(", numQueries=").append(numQueries);
sb.append('}');
return sb.toString();
}

public String getToken() {
return token;
}

public Quota setToken(String token) {
this.token = token;
return this;
}

public String getDate() {
return date;
}

public Quota setDate(int year, int month) {
this.date = year + String.format("%02d", month);
return this;
}

public Quota setDate(String date) {
this.date = date;
return this;
}

public long getNumQueries() {
return numQueries;
}

public Quota setNumQueries(long numQueries) {
this.numQueries = numQueries;
return this;
}
}
Loading

0 comments on commit 8dbb33f

Please sign in to comment.