ucl/ucl/builtins/bytes.go

package builtins

import (
	"bytes"
	"context"
	"crypto/md5"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"

	"ucl.lmika.dev/ucl"
)

type bytesHandlers struct{}

/// :module bytes
//
// Perform operations over byte slices.

func Bytes() ucl.Module {
	bh := bytesHandlers{}

	return ucl.Module{
		Name: "bytes",
		Builtins: map[string]ucl.BuiltinHandler{
			"from":   bh.from,
			"hash":   bh.hash,
			"base64": bh.base64,
			"hex":    bh.hex,
		},
	}
}

// :fn hash
//
// :syntax INPUT ALGORITHM
//
// Returns a byte slice containing the result of hashing the input with the given algorithm.
//
// The input can either be a byte slice or a string. The supported algorithms are as follows:
//
// - md5
// - sha1
// - sha256
func (b bytesHandlers) hash(ctx context.Context, args ucl.CallArgs) (any, error) {
	inputBts, args, err := consumeStringOrBytes(args)
	if err != nil {
		return nil, err
	}

	var algor string
	if err := args.Bind(&algor); err != nil {
		return nil, err
	}

	switch algor {
	case "md5":
		res := md5.Sum(inputBts)
		return ucl.BytesObject(res[:]), nil
	case "sha1":
		res := sha1.Sum(inputBts)
		return ucl.BytesObject(res[:]), nil
	case "sha256":
		res := sha256.Sum256(inputBts)
		return ucl.BytesObject(res[:]), nil
	}

	return nil, errors.New("unsupported algorithm")
}

// :fn base64
//
// :syntax INPUT
//
// Returns a string containing the base64 encoding of the input.
//
// The input can either be a byte slice or a string. The result will be a base64 encoded string
// using standard encoding.
func (c bytesHandlers) base64(ctx context.Context, args ucl.CallArgs) (any, error) {
	inputBts, args, err := consumeStringOrBytes(args)
	if err != nil {
		return nil, err
	}

	res := base64.StdEncoding.EncodeToString(inputBts)
	return res, nil
}

// :fn from
//
// :syntax INPUT ...
//
// Returns a byte slice from a given input. Input can be one of the following:
//
// - Nil, which will produce an empty byte slice
// - A string, which would produce a byte slice containing the UTF-8 encoded string
// - A number, which would produce a single byte slice containing the number as a byte
// - A byte slice, which would produce a copy of the byte slice
//
// A list or iter will consume the elements and apply the byte slice conversion recursively, combining
// the results into a single byte slice.
func (c bytesHandlers) from(ctx context.Context, args ucl.CallArgs) (any, error) {
	var bfr bytes.Buffer

	var o ucl.Object
	for args.NArgs() > 0 {
		if err := args.Bind(&o); err != nil {
			return nil, fmt.Errorf("failed to bind object: %w", err)
		}
		if err := writeObjToBytesBuffer(ctx, &bfr, o); err != nil {
			return nil, err
		}
	}

	return bfr.Bytes(), nil
}

// :fn hex
//
// :syntax BYTES
//
// Returns a string encoding the bytes slice as a hex string
func (c bytesHandlers) hex(ctx context.Context, args ucl.CallArgs) (any, error) {
	inputBts, args, err := consumeStringOrBytes(args)
	if err != nil {
		return nil, err
	}

	res := hex.EncodeToString(inputBts)
	return res, nil
}

func consumeStringOrBytes(args ucl.CallArgs) (bs ucl.BytesObject, _ ucl.CallArgs, _ error) {
	if objs := args.RestAsObjects(); len(objs) > 0 {
		if b, ok := objs[0].(ucl.BytesObject); ok {
			args.Shift(1)
			return b, args, nil
		}
	}

	var s string
	if err := args.Bind(&s); err != nil {
		return nil, ucl.CallArgs{}, err
	}
	return ucl.BytesObject(s), args, nil
}

func writeObjToBytesBuffer(ctx context.Context, bfr *bytes.Buffer, o ucl.Object) error {
	if o == nil {
		return nil
	}

	switch v := o.(type) {
	case ucl.StringObject:
		bfr.Write([]byte(v))
	case ucl.IntObject:
		bfr.WriteByte(byte(v))
	case ucl.StringListObject:
		for _, el := range v {
			bfr.Write([]byte(el))
		}
	case ucl.BytesObject:
		bfr.Write(v)
	case ucl.Listable:
		l := v.Len()
		for i := 0; i < l; i++ {
			if err := writeObjToBytesBuffer(ctx, bfr, v.Index(i)); err != nil {
				return err
			}
		}
	case ucl.Iterable:
		for v.HasNext() {
			n, err := v.Next(ctx)
			if err != nil {
				return err
			}
			if n == nil {
				break
			}
			if err := writeObjToBytesBuffer(ctx, bfr, n); err != nil {
				return err
			}
		}
	default:
		return fmt.Errorf("unsupported type %T", o)
	}
	return nil
}