<commit_before>// Simple utility for counting messages on a Kafka topic.
//
// Copyright (C) 2017 ENEO Tecnologia SL
// Author: <NAME> <<EMAIL>>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"strings"
	"time"

	"github.com/fatih/color"
)

// PubKey should be set on compile time
var PubKey string

// License contains information about the limits of bytes that an organization
// can send.
type License struct {
	UUID        string         `json:"uuid"`
	ExpireAt    int64          `json:"expire_at"`
	LimitBytes  int64          `json:"limit_bytes"`
	ClusterUUID string         `json:"cluster_uuid"`
	Sensors     map[string]int `json:"sensors"`
}

// LoadLicenses reads a redBorder license from a file and returns a License
// struct holding the decoded information.
func LoadLicenses(config *AppConfig) (int64, error) {
	var licenses []License
	log := log.WithField("prefix", "license")

	files, err := ioutil.ReadDir(config.LicensesDirectory)
	if err != nil {
		return -1, err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		encodedLicense, encodedSignature, err :=
			ReadLicense(config.LicensesDirectory + "/" + file.Name())
		if err != nil {
			return -1, err
		}

		if PubKey != "" {
			pub, err := DecodePEM([]byte(PubKey))
			if err != nil {
				return -1, err
			}

			valid, err := VerifyLicense(pub, encodedLicense, encodedSignature)
			if err != nil {
				return -1, err
			}

			if !valid {
				return -1, errors.New("Signature not valid")
			}
		}

		license, err := DecodeLicense(encodedLicense)
		if err != nil {
			return -1, err
		}

		log.Infoln(FormatLicense(license))

		licenses = append(licenses, *license)
	}

	var totalBytes int64
	for _, license := range licenses {
		totalBytes += license.LimitBytes
	}

	log.Infof(color.GreenString("Total bytes allowed: %d", totalBytes))

	return totalBytes, nil
}

// ReadLicense read the license from a file. Decode as a base64 string and
// parses it as a JSON.
func ReadLicense(filename string) (encodedInfo, signature []byte, err error) {
	var (
		encodedSignatureStr string
		encodedInfoStr      string
		ok                  bool
	)

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}

	var el map[string]interface{}
	err = json.Unmarshal(data, &el)
	if err != nil {
		return nil, nil, err
	}

	if encodedInfoStr, ok = el["encoded_info"].(string); !ok {
		return nil, nil, errors.New("Encoded info is not string")
	}

	if encodedSignatureStr, ok = el["signature"].(string); !ok {
		return nil, nil, errors.New("Signature info is not string")
	}

	return []byte(encodedInfoStr), []byte(encodedSignatureStr), nil
}

// DecodePEM reads and parses a PEM key with a public key
func DecodePEM(key []byte) (*rsa.PublicKey, error) {
	pemKey := strings.Replace(string(key), `\n`, "\n", -1)
	block, _ := pem.Decode([]byte(pemKey))
	if block == nil {
		return nil, errors.New("Error decoding PEM")
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	pubKey, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("Key is not an RSA key")
	}

	return pubKey, nil
}

// VerifyLicense checks if a encoded license has a valid signature using a
// public key.
func VerifyLicense(
	pub *rsa.PublicKey, encodedLicense, encodedSignature []byte,
) (bool, error) {
	var err error

	hashed := sha256.Sum256(encodedLicense)
	signatureLen := base64.URLEncoding.DecodedLen(len(encodedSignature))
	signature := make([]byte, signatureLen)
	n, err := base64.URLEncoding.Decode(signature, encodedSignature)
	if err != nil {
		return false, err
	}

	err = rsa.VerifyPKCS1v15(pub, crypto.SHA256, hashed[:], signature[:n])
	if err != nil {
		return false, err
	}

	return true, nil
}

// DecodeLicense receives a JSON marshaled license a returns a struct with the
// information.
func DecodeLicense(encodedLicense []byte) (*License, error) {
	license := new(License)

	dataLen := base64.URLEncoding.DecodedLen(len(encodedLicense))
	data := make([]byte, dataLen)

	n, err := base64.URLEncoding.Decode(data, encodedLicense)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(data[:n], license)
	if err != nil {
		return nil, err
	}

	return license, nil
}

// FormatLicense create a pretty string given a license object
func FormatLicense(license *License) string {<commit_msg>:wrench: Improve license verifying<commit_after>// Simple utility for counting messages on a Kafka topic.
//
// Copyright (C) 2017 ENEO Tecnologia SL
// Author: <NAME> <<EMAIL>>
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

package main

import (
	"crypto"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/json"
	"encoding/pem"
	"errors"
	"fmt"
	"io/ioutil"
	"strings"
	"time"

	"github.com/fatih/color"
)

// PubKey should be set on compile time
var PubKey string

// License contains information about the limits of bytes that an organization
// can send.
type License struct {
	UUID        string         `json:"uuid"`
	ExpireAt    int64          `json:"expire_at"`
	LimitBytes  int64          `json:"limit_bytes"`
	ClusterUUID string         `json:"cluster_uuid"`
	Sensors     map[string]int `json:"sensors"`
}

// LoadLicenses reads a redBorder license from a file and returns a License
// struct holding the decoded information.
func LoadLicenses(config *AppConfig) (int64, error) {
	var licenses []License
	log := log.WithField("prefix", "license")

	files, err := ioutil.ReadDir(config.LicensesDirectory)
	if err != nil {
		return -1, err
	}

	for _, file := range files {
		if file.IsDir() {
			continue
		}

		encodedLicense, encodedSignature, err :=
			ReadLicense(config.LicensesDirectory + "/" + file.Name())
		if err != nil {
			return -1, err
		}

		if PubKey != "" {
			pub, err := DecodePEM([]byte(PubKey))
			if err != nil {
				return -1, err
			}

			err = VerifyLicense(pub, encodedLicense, encodedSignature)
			if err != nil {
				log.Errorf(
					"Error verifying license \"%s\": %s", file.Name(), err.Error())
				continue
			}
		}

		license, err := DecodeLicense(encodedLicense)
		if err != nil {
			return -1, err
		}

		expires := time.Unix(license.ExpireAt, 0)
		if expires.Before(time.Now()) {
			log.Warnf("License %s has expired", license.UUID)
			continue
		}

		log.Infoln(FormatLicense(license))

		licenses = append(licenses, *license)
	}

	var totalBytes int64
	for _, license := range licenses {
		totalBytes += license.LimitBytes
	}

	log.Infof(color.GreenString("Total bytes allowed: %d", totalBytes))

	return totalBytes, nil
}

// ReadLicense read the license from a file. Decode as a base64 string and
// parses it as a JSON.
func ReadLicense(filename string) (encodedInfo, signature []byte, err error) {
	var (
		encodedSignatureStr string
		encodedInfoStr      string
		ok                  bool
	)

	data, err := ioutil.ReadFile(filename)
	if err != nil {
		return nil, nil, err
	}

	var el map[string]interface{}
	err = json.Unmarshal(data, &el)
	if err != nil {
		return nil, nil, err
	}

	if encodedInfoStr, ok = el["encoded_info"].(string); !ok {
		return nil, nil, errors.New("Encoded info is not string")
	}

	if encodedSignatureStr, ok = el["signature"].(string); !ok {
		return nil, nil, errors.New("Signature info is not string")
	}

	return []byte(encodedInfoStr), []byte(encodedSignatureStr), nil
}

// DecodePEM reads and parses a PEM key with a public key
func DecodePEM(key []byte) (*rsa.PublicKey, error) {
	pemKey := strings.Replace(string(key), `\n`, "\n", -1)
	block, _ := pem.Decode([]byte(pemKey))
	if block == nil {
		return nil, errors.New("Error decoding PEM")
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, err
	}

	pubKey, ok := pub.(*rsa.PublicKey)
	if !ok {
		return nil, errors.New("Key is not an RSA key")
	}

	return pubKey, nil
}

// VerifyLicense checks if a encoded license has a valid signature using a
// public key.
func VerifyLicense(
	pub *rsa.PublicKey, encodedLicense, encodedSignature []byte,
) error {
	hashed := sha256.Sum256(encodedLicense)
	signatureLen := base64.URLEncoding.DecodedLen(len(encodedSignature))
	signature := make([]byte, signatureLen)

	n, err := base64.URLEncoding.Decode(signature, encodedSignature)
	if err != nil {
		return err
	}

	err = rsa.VerifyPKCS1v15(pub, crypto.SHA256, hashed[:], signature[:n])
	if err != nil {
		return err
	}

	return nil
}

// DecodeLicense receives a JSON marshaled license a returns a struct with the
// information.
func DecodeLicense(encodedLicense []byte) (*License, error) {
	license := new(License)

	dataLen := base64.URLEncoding.DecodedLen(len(encodedLicense))
	data := make([]byte, dataLen)

	n, err := base64.URLEncoding.Decode(data, encodedLicense)
	if err != nil {
		return nil, err
	}

	err = json.Unmarshal(data[:n], license)
	if err != nil {
		return nil, err
	}

	return license, nil
}

// FormatLicense create a pretty string given a license object
func FormatLicense(license *License) string {
{'fraction_non_alphanumeric': 0.12081120231772091, 'fraction_numerical': 0.010140028971511348, 'mean_word_length': 3.5802742149491373, 'pattern_counts': {'":': 1, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5781010', 'n_tokens_mistral': 3555, 'n_tokens_neox': 3293, 'n_words': 1107}