Google CTF 2021 Tonality

// Copyright 2021 Google LLC
//
// 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 main

import (
    "crypto-tonality/challenger"
    "crypto-tonality/pb"
    "encoding/binary"
    "flag"
    "fmt"
    "io"
    "io/ioutil"
    "os"
    "strings"

    "github.com/golang/protobuf/proto"
)

var flagFile string

const maxMessageLengh = 10 * 1024 * 1024

func writeMessage(w io.Writer, m proto.Message) error {
    if err := binary.Write(w, binary.LittleEndian, uint32(proto.Size(m))); err != nil {
        return fmt.Errorf("failed to write message length, binary.Write() = %v, want nil err", err)
    }
    buf, err := proto.Marshal(m)
    if err != nil {
        return fmt.Errorf("failed to serialize message, proto.Marshal(%v) = %v, want nil err", m, err)
    }
    n, err := w.Write(buf)
    if n != len(buf) || err != nil {
        return fmt.Errorf("failed to write serialized message, w.Write(%X) = %v, %v, want n=%d and nil error", buf, n, err, len(buf))
    }
    return nil
}

func readMessage(r io.Reader, m proto.Message) error {
    var length uint32
    if err := binary.Read(r, binary.LittleEndian, &length); err != nil {
        return fmt.Errorf("failed to read length, binary.Read(buf) = %v, want nil error", err)
    }
    if length > maxMessageLengh {
        return fmt.Errorf("want legnth <= %d, got %d", maxMessageLengh, length)
    }
    buf := make([]byte, length)
    if n, err := r.Read(buf); n != len(buf) || err != nil {
        return fmt.Errorf("failed to read serialized message, r.Read(buf) = %v, %v, want n=%d and nil error", n, err, len(buf))
    }
    if err := proto.Unmarshal(buf, m); err != nil {
        return fmt.Errorf("failed to unmarshal message, proto.Unmarshal(%X, m) = %v, want nil error", buf, err)
    }
    return nil
}

func runSession(chal *challenger.Challenger, r io.Reader, w io.Writer) error {
    hello := chal.Hello(&pb.HelloRequest{})
    if err := writeMessage(w, hello); err != nil {
        return fmt.Errorf("writeMessage(w, hello=%X)=%v, want nil err", hello, err)
    }

    signReq := &pb.SignRequest{}
    if err := readMessage(r, signReq); err != nil {
        return fmt.Errorf("readMessage(r, signReq)=%v, want nil err", err)
    }

    signRes := chal.SignFirstMessage(signReq)
    if err := writeMessage(w, signRes); err != nil {
        return fmt.Errorf("writeMessage(w, signRes=%X)=%v, want nil err", signRes, err)
    }

    verifyReq := &pb.VerifyRequest{}
    if err := readMessage(r, verifyReq); err != nil {
        return fmt.Errorf("readMessage(r, verifyReq)=%v, want nil err", err)
    }

    verifyRes := chal.VerifySecondMessage(verifyReq)
    if err := writeMessage(w, verifyRes); err != nil {
        return fmt.Errorf("writeMessage(w, verifyRes=%X)=%v, want nil err", verifyRes, err)
    }
    return nil
}

func init() {
    flag.StringVar(&flagFile, "flag", "/flag", "flag filename")
}

func main() {
    defer func() {
        if r := recover(); r != nil {
            // Uncomment to catch panics during development.
            // panic(r)
        }
    }()

    flag.Parse()

    f, err := ioutil.ReadFile(flagFile)
    if err != nil {
        panic(err)
    }

    chal, err := challenger.NewChallenger(strings.TrimSpace(string(f)))
    if err != nil {
        panic(err)
    }

    err = runSession(chal, os.Stdin, os.Stdout)
    if err != nil {
        panic(err)
    }
}

// Copyright 2021 Google LLC
//
// 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 challenger

import (
    "crypto-tonality/pb"
    "crypto/ecdsa"
    "crypto/elliptic"
    "crypto/rand"
    "crypto/sha1"
    "io"
    "math/big"
)

const m0 = "Server says 1+1=2"
const m1 = "Server says 1+1=3"

func hashMessage(m string) []byte {
    h := sha1.New()
    io.WriteString(h, m)
    return h.Sum(nil)
}

type Challenger struct {
    flag string
    sk   *ecdsa.PrivateKey
}

func NewChallenger(flag string) (*Challenger, error) {
    sk, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
    if err != nil {
        return nil, err
    }
    return &Challenger{
        flag: flag,
        sk:   sk,
    }, nil
}

func (chal *Challenger) Hello(req *pb.HelloRequest) *pb.HelloResponse {
    return &pb.HelloResponse{
        Message0: m0,
        Message1: m1,
        Pubkey: &pb.Point{
            X: chal.sk.PublicKey.X.Bytes(),
            Y: chal.sk.PublicKey.Y.Bytes(),
        },
    }
}

func (chal *Challenger) scalePrivate(t *big.Int) *ecdsa.PrivateKey {
    rk := new(big.Int).Mul(chal.sk.D, t)
    rk.Mod(rk, chal.sk.PublicKey.Curve.Params().N)

    ret := new(ecdsa.PrivateKey)
    ret.PublicKey.Curve = chal.sk.PublicKey.Curve
    ret.D = rk
    ret.PublicKey.X, ret.PublicKey.Y = chal.sk.PublicKey.Curve.ScalarBaseMult(ret.D.Bytes())
    return ret
}

func (chal *Challenger) SignFirstMessage(req *pb.SignRequest) *pb.SignResponse {
    t := new(big.Int).SetBytes(req.Scalar)
    r, s, err := ecdsa.Sign(rand.Reader, chal.scalePrivate(t), hashMessage(m0))
    if err != nil {
        return &pb.SignResponse{}
    }
    return &pb.SignResponse{
        Message0Sig: &pb.Signature{
            R: r.Bytes(),
            S: s.Bytes(),
        },
    }
}

func (chal *Challenger) VerifySecondMessage(req *pb.VerifyRequest) *pb.VerifyResponse {
    r := new(big.Int).SetBytes(req.Message1Sig.R)
    s := new(big.Int).SetBytes(req.Message1Sig.S)

    ok := ecdsa.Verify(&chal.sk.PublicKey, hashMessage(m1), r, s)
    if !ok {
        return &pb.VerifyResponse{}
    }

    return &pb.VerifyResponse{Flag: chal.flag}
}

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2021 Google LLC
#
# 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
#
#     https://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.

import argparse
import pwnlib
import challenge_pb2
import struct
import sys

def handle_pow(tube):
  raise NotImplemented()

def read_message(tube, typ):
  n = struct.unpack('<L', tube.recvnb(4))[0]
  buf = tube.recvnb(n)
  msg = typ()
  msg.ParseFromString(buf)
  return msg

def write_message(tube, msg):
  buf = msg.SerializeToString()
  tube.send(struct.pack('<L', len(buf)))
  tube.send(buf)

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--port', metavar='P', type=int, default=1337, help='challenge #port')
  parser.add_argument(
      '--host', metavar='H', type=str, default='tonality.2021.ctfcompetition.com', help='challenge host')
  args = parser.parse_args()

  tube = pwnlib.tubes.remote.remote(args.host, args.port)
  print(tube.recvuntil('== proof-of-work: '))
  if tube.recvline().startswith(b'enabled'):
      handle_pow(tube)

  # Step 1: Hello.
  hello = read_message(tube, challenge_pb2.HelloResponse)
  print(hello)

  # Step 2: Sign.
  a = 1234
  sign_req = challenge_pb2.SignRequest()
  sign_req.scalar = a.to_bytes((a.bit_length() + 7) // 8, 'big')
  write_message(tube, sign_req)

  sign_res = read_message(tube, challenge_pb2.SignResponse)
  print(sign_res)

  # Step 3: Verify.
  verify_req = challenge_pb2.VerifyRequest()
  verify_req.message1_sig.r = b'\x11\x22'
  verify_req.message1_sig.s = b'\x33\x44'
  write_message(tube, verify_req)

  verify_res = read_message(tube, challenge_pb2.VerifyResponse)
  print(verify_res)
  return 0


if __name__ == '__main__':
  sys.exit(main())