enigma.go

package main

import (
	"errors"
	"fmt"
	"log"
	"regexp"
	"strings"
)

const (
	alphabet     = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	rotorConfI   = "EKMFLGDQVZNTOWYHXUSPAIBRCJ"
	rotorConfII  = "AJDKSIRUXBLHWTMCQGZNPYFVOE"
	rotorConfIII = "BDFHJLCPRTXVZNYEIWGAKMUSQO"
	rotorConfIV  = "ESOVPZJAYQUIRHXLNFTGKDCMWB"
	rotorConfV   = "VZBRGITYUPSDNHLXAWMJQOFECK"
	reflector    = "EJMZALYXVBWFCRQUONTSPIKHGD"
)

type Rotor struct {
	letters       string
	position      int
	notchPosition int
}

func (r *Rotor) substituteForward(letter byte) byte {
	return r.letters[(strings.IndexByte(alphabet, letter)+r.position)%26]
}

func (r *Rotor) substituteReverse(letter byte) byte {
	i := strings.IndexByte(r.letters, letter)
	i = i - r.position
	if i < 0 {
		i = 26 + i
	}
	return alphabet[i]
}

var (
	RotorI = Rotor{
		letters:       rotorConfI,
		notchPosition: 8,
	}
	RotorII = Rotor{
		letters:       rotorConfII,
		notchPosition: 0,
	}
	RotorIII = Rotor{
		letters:       rotorConfIII,
		notchPosition: 12,
	}
	RotorIV = Rotor{
		letters:       rotorConfIV,
		notchPosition: 7,
	}
	RotorV = Rotor{
		letters:       rotorConfV,
		notchPosition: 2,
	}
)

type Configuration struct {
	Rotors            [3]Rotor
	RingSettings      [3]int
	StartingPositions [3]int
	PlugboardConf     map[byte]byte
}

type Enigma struct {
	rotors    [3]Rotor
	plugBoard map[byte]byte
}

func NewEnigma(conf Configuration) (*Enigma, error) {
	reg, err := regexp.Compile("[^a-zA-Z]+")
	if err != nil {
		return nil, err
	}
	if conf.PlugboardConf == nil {
		conf.PlugboardConf = make(map[byte]byte)
	}
	if len(conf.PlugboardConf) > 13 {
		return nil, errors.New("plugboard cannot have more than 13 mappings")
	}
	for k, v := range conf.PlugboardConf {
		if reg.Match([]byte{k}) || reg.Match([]byte{v}) {
			return nil, errors.New("plugboard contains invalid characters")
		}
		conf.PlugboardConf[v] = k
	}

	for i := 0; i < 3; i++ {
		if conf.RingSettings[i] < 0 || conf.RingSettings[i] > 25 {
			return nil, errors.New("ring settings must be between 0 and 25")
		}
		if conf.StartingPositions[i] < 0 || conf.StartingPositions[i] > 25 {
			return nil, errors.New("starting positions must be between 0 and 25")
		}
	}
	e := &Enigma{
		rotors:    conf.Rotors,
		plugBoard: conf.PlugboardConf,
	}
	e.rotors[0].notchPosition = (e.rotors[0].notchPosition + conf.RingSettings[0]) % 26
	e.rotors[1].notchPosition = (e.rotors[1].notchPosition + conf.RingSettings[1]) % 26
	e.rotors[2].notchPosition = (e.rotors[2].notchPosition + conf.RingSettings[2]) % 26

	for i := 0; i < 3; i++ {
		for e.rotors[i].position != conf.StartingPositions[i] {
			e.rotate(i)
		}
	}
	return e, nil
}

func (e *Enigma) rotate(rotor int) {
	if rotor < 0 || rotor > 2 {
		panic("Rotor index out of bounds")
	}
	if rotor < 2 && e.rotors[rotor].position == e.rotors[rotor].notchPosition {
		e.rotate(rotor + 1)
	}
	e.rotors[rotor].position = (e.rotors[rotor].position + 1) % 26
}

func (e *Enigma) rotorSubstitute(letter byte) byte {
	e.rotate(0)
	l1 := e.rotors[0].substituteForward(letter)
	l2 := e.rotors[1].substituteForward(l1)
	l3 := e.rotors[2].substituteForward(l2)
	l4 := reflector[strings.IndexByte(alphabet, l3)]
	l5 := e.rotors[2].substituteReverse(l4)
	l6 := e.rotors[1].substituteReverse(l5)
	output := e.rotors[0].substituteReverse(l6)
	return output
}

func (e *Enigma) plugboardSubstitute(letter byte) byte {
	sub, ok := e.plugBoard[letter]
	if !ok {
		return letter
	}
	return sub
}

func (e *Enigma) Encrypt(plainText string) string {
	reg, err := regexp.Compile("[^a-zA-Z]+")
	if err != nil {
		panic("regex error")
	}
	plainText = strings.ToUpper(reg.ReplaceAllString(plainText, ""))
	var cipherText string
	for i := 0; i < len(plainText); i++ {
		sub0 := e.plugboardSubstitute(plainText[i])
		sub1 := e.rotorSubstitute(sub0)
		out := e.plugboardSubstitute(sub1)
		cipherText += string(out)
	}
	return cipherText
}

func (e *Enigma) Decrypt(cipherText string) string {
	var plainText string
	for i := 0; i < len(cipherText); i++ {
		sub0 := e.plugboardSubstitute(cipherText[i])
		sub1 := e.rotorSubstitute(sub0)
		out := e.plugboardSubstitute(sub1)
		plainText += string(out)
	}
	return plainText
}

func main() {
	conf := Configuration{
		Rotors:            [3]Rotor{RotorI, RotorII, RotorIII},
		RingSettings:      [3]int{17, 21, 3},
		StartingPositions: [3]int{5, 3, 11},
		PlugboardConf: map[byte]byte{
			'A': 'J',
			'B': 'R',
			'Z': 'D',
			'M': 'O',
		},
	}
	enigma1, err := NewEnigma(conf)
	if err != nil {
		log.Fatal(err)
	}

	cipherText := enigma1.Encrypt("Hello World")
	fmt.Println("Ciphertext:", cipherText)

	enigma2, err := NewEnigma(conf)
	if err != nil {
		log.Fatal(err)
	}

	plaintext := enigma2.Decrypt(cipherText)
	fmt.Println("Plaintext:", plaintext)
}