package main

import (
  "bytes"
  "crypto/md5"
  "encoding/binary"
  "errors"
  "fmt"
  "io"
  "strconv"
)

const MT_BOOL_FALSE byte = 0x00
const MT_BOOL_TRUE byte = 0x01
const MT_DWORD byte = 0x08
const MT_BYTE byte = 0x09
const MT_STRING byte = 0x21
const MT_HASH byte = 0x31
const MT_ARRAY byte = 0x88

const MT_RECEIVER = 0xFF0001
const MT_SENDER = 0xFF0002
const MT_REPLY_EXPECTED = 0xFF0005
const MT_REQUEST_ID = 0xFF0006
const MT_COMMAND = 0xFF0007

type M2Element struct {
  code int
  value interface{}
}

func (el *M2Element) String() string {
  return fmt.Sprintf("code=%v,type=%T,value=%v", el.code, el.value, el.value)
}


type M2Message struct {
  el []M2Element
}

type M2Hash string

func NewM2Message() *M2Message {
  m2 := M2Message{}
  return &m2
}

func (m2 *M2Message) Clear() {
  m2.el = []M2Element{}
}

func (m2 *M2Message) Append(code int, value interface{}) {
  if m2.el == nil {
    m2.Clear()
  }
  m2.el = append(m2.el, M2Element{code: code, value: value})
}

func (m2 *M2Message) AppendElement(el *M2Element) {
  if m2.el == nil {
    m2.Clear()
  }
  m2.el = append(m2.el, *el)
}

func (m2 *M2Message) Bytes() []byte {
  res := []byte{}

  for _, el := range m2.el {
    buf := new(bytes.Buffer)

    binary.Write(buf, binary.LittleEndian, uint16(el.code))
    binary.Write(buf, binary.LittleEndian, byte(el.code >> 16))

    switch v := el.value.(type) {
    case bool:
      binary.Write(buf, binary.LittleEndian, v)
    case byte:
      binary.Write(buf, binary.LittleEndian, byte(MT_BYTE))
      binary.Write(buf, binary.LittleEndian, v)
    case int:
      binary.Write(buf, binary.LittleEndian, byte(MT_DWORD))
      binary.Write(buf, binary.LittleEndian, int32(v))
    case uint:
      binary.Write(buf, binary.LittleEndian, byte(MT_DWORD))
      binary.Write(buf, binary.LittleEndian, uint32(v))
    case string:
      binary.Write(buf, binary.LittleEndian, byte(MT_STRING))
      binary.Write(buf, binary.LittleEndian, byte(len(v)))
      binary.Write(buf, binary.LittleEndian, []byte(v))
    case M2Hash:
      binary.Write(buf, binary.LittleEndian, byte(MT_HASH))
      binary.Write(buf, binary.LittleEndian, byte(len(v)))
      binary.Write(buf, binary.LittleEndian, []byte(v))
    case []byte:
      binary.Write(buf, binary.LittleEndian, byte(MT_ARRAY))
      binary.Write(buf, binary.LittleEndian, uint16(len(v)))
      for _, i := range v {
        binary.Write(buf, binary.LittleEndian, int32(i))
      }
    case []int:
      binary.Write(buf, binary.LittleEndian, byte(MT_ARRAY))
      binary.Write(buf, binary.LittleEndian, uint16(len(v)))
      for _, i := range v {
        binary.Write(buf, binary.LittleEndian, int32(i))
      }
    }

    res = append(res, buf.Bytes()...)
  }

  header := make([]byte, 6)
  header[0] = byte(len(res) + 4)
  header[1] = 0x01
  header[2] = 0x00
  header[3] = byte(len(res) + 2)
  header[4] = 0x4D
  header[5] = 0x32

  return append(header, res...)
}


func (m2 *M2Message) ParseM2Element(buf io.Reader) error {
  var codeAndType uint32
  err := binary.Read(buf, binary.LittleEndian, &codeAndType)
  if err != nil {
    return err
  }

  el := M2Element{code: int(codeAndType & 0x00FFFFFF)}
  keyType := byte(codeAndType >> 24)
  log("lw", 3, "m2 code=%v type=%v", el.code, keyType)

  switch keyType {
  case MT_BOOL_FALSE, MT_BOOL_TRUE:
    el.value = keyType == MT_BOOL_TRUE
    log("lw", 3, "m2 MT_BOOL: %v", el.value.(bool))
  case MT_BYTE:
    var b byte
    err = binary.Read(buf, binary.LittleEndian, &b)
    el.value = b
    log("lw", 3, "m2 MT_BYTE: %v", el.value.(byte))
  case MT_DWORD:
    var b int32
    err = binary.Read(buf, binary.LittleEndian, &b)
    el.value = b
    log("lw", 3, "m2 MT_DWORD: %v", el.value.(int32))

  case MT_STRING:
    var length byte
    err = binary.Read(buf, binary.LittleEndian, &length)
    if err != nil {
      return err
    }

    bs := make([]byte, length)
    _, err = io.ReadFull(buf, bs)
    el.value = string(bs)
    log("lw", 3, "m2 MT_STRING (len %v): %v", length, el.value.(string))

  case MT_HASH:
    var length byte
    err = binary.Read(buf, binary.LittleEndian, &length)
    if err != nil {
      return err
    }

    bs := make([]byte, length)
    _, err = io.ReadFull(buf, bs)
    el.value = M2Hash(bs)
    log("lw", 3, "m2 MT_HASH (len %v): %v", length, []byte(el.value.(M2Hash)))

  case MT_ARRAY:
    var length uint16
    err = binary.Read(buf, binary.LittleEndian, &length)
    if err != nil {
      return err
    }

    sl := []int{}
    for i := 0; i < int(length); i++ {
      var el2 int32
      err = binary.Read(buf, binary.LittleEndian, &el2)
      if err != nil {
        break
      }
      sl = append(sl, int(el2))
    }
    el.value = sl
    log("lw", 3, "m2 MT_HASH (len %v): %v", length, el.value.([]int))

  default:
    return errors.New("unknown key code " + strconv.Itoa(int(keyType)))
  }

  if err != nil {
    return err
  }

  m2.el = append(m2.el, el)
  return nil
}

func (m2 *M2Message) ParseM2Message(buf io.Reader) error {
  var headerBlockSize, m2BlockSize byte
  var m2Extra, m2Header uint16
  err := binary.Read(buf, binary.LittleEndian, &headerBlockSize)
  err = binary.Read(buf, binary.LittleEndian, &m2Extra)
  err = binary.Read(buf, binary.LittleEndian, &m2BlockSize)
  err = binary.Read(buf, binary.LittleEndian, &m2Header)
  if err != nil {
    return err
  }
  if m2Extra != 0x1 {
    return errors.New("invalid M2_EXTRA")
  }
  if m2Header != 0x324D {
    return errors.New("invalid M2_HEADER")
  }

  for {
    log("lw", 3, "parsing new m2 element")
    err := m2.ParseM2Element(buf)
    if err != nil {
      return err
    }
  }
}

func ParseM2Messages(src []byte) (messages []M2Message, err error) {
  messages = []M2Message{}
	buf := bytes.NewReader(src)

  for {
    m2 := NewM2Message()
    err := m2.ParseM2Message(buf)

    if err == io.EOF {
      messages = append(messages, *m2)
      break
    } else if err != nil {
      return nil, err
    } else {
      messages = append(messages, *m2)
    }
  }


  log("lw", 3, "m2 eof after %v messages", len(messages))
  return messages, nil
}










type LegacyWinbox struct {
  task *Task
  stage int
  m2 []M2Message
}

func NewLegacyWinbox(task *Task) *LegacyWinbox {
  lw := LegacyWinbox{task: task, stage: -1, m2: []M2Message{}}
  return &lw
}


// req1
func (lw *LegacyWinbox) MTReqList() []byte {
  m2 := NewM2Message()
  m2.Append(MT_RECEIVER, []byte{2, 2})
  m2.Append(MT_COMMAND, byte(7))
  m2.Append(MT_REQUEST_ID, byte(1))
  m2.Append(MT_REPLY_EXPECTED, true)
  m2.Append(1, "list")
  return m2.Bytes()
}

// res1
func (lw *LegacyWinbox) MTGetSid(m2 []M2Message) *M2Element {
  for _, msg := range m2 {
    for _, el := range msg.el {
      if el.code == 0xFE0001 {
        return &el
      }
    }
  }

  return nil
}


// req2
func (lw *LegacyWinbox) MTReqChallenge(sid *M2Element) []byte {
  m2 := NewM2Message()
  m2.Append(MT_RECEIVER, []byte{13, 4})
  m2.Append(MT_COMMAND, byte(4))
  m2.Append(MT_REQUEST_ID, byte(2))
  m2.AppendElement(sid)
  m2.Append(MT_REPLY_EXPECTED, true)
  return m2.Bytes()
}

// res2
func (lw *LegacyWinbox) MTGetSalt(m2 []M2Message) M2Hash {
  for _, msg := range m2 {
    for _, el := range msg.el {
      if el.code == 0x9 {
        return el.value.(M2Hash)
      }
    }
  }

  return ""
}


// req3
func (lw *LegacyWinbox) MTReqAuth(sid *M2Element, login, digest, salt string) []byte {
  m2 := NewM2Message()
  m2.Append(MT_RECEIVER, []byte{13, 4})
  m2.Append(MT_COMMAND, byte(1))
  m2.Append(MT_REQUEST_ID, byte(3))
  m2.AppendElement(sid)
  m2.Append(MT_REPLY_EXPECTED, true)
  m2.Append(1, login)
  m2.Append(9, M2Hash(salt))
  m2.Append(10, M2Hash(digest))
  return m2.Bytes()
}

// res3
func (lw *LegacyWinbox) MTGetResult(m2 []M2Message) (res bool, err error) {
  for _, msg := range m2 {
    for _, el := range msg.el {
      if el.code == 0xA {
        _, ok := el.value.(M2Hash)
        if ok {
          return true, nil
        }
      } else if el.code == 0xFF0008 {
        v, ok := el.value.(int32)
        if ok && v == 0xFE0006 {
          return false, nil
        }
      }
    }
  }

  return false, errors.New("no auth marker found")
}




func (lw *LegacyWinbox) SendRecv(buf []byte) (res []byte, err error) {
  _, err = lw.task.conn.Write(buf)
  if err != nil {
    log("lw", 1, "failed to send: %v", err.Error())
    return nil, err
  }

  resp := make([]byte, 1024)
  n, err := lw.task.conn.Read(resp)
  if err != nil {
    log("lw", 1, "failed to recv: %v", err.Error())
    return nil, err
  }

  return resp[:n], nil
}


func (lw *LegacyWinbox) TryLogin() (res bool, err error) {
  log("lw", 2, "login: stage 1, req_list")
  r1, err := lw.SendRecv(lw.MTReqList())
  if err != nil {
    return false, err
  }

  log("lw", 2, "login: stage 2, got response for req_list")
  msg, err := ParseM2Messages(r1)
  if err != nil {
    return false, err
  }

  sid := lw.MTGetSid(msg)
  if sid == nil {
    return false, errors.New("failed to get SID from stage 2")
  }
  log("lw", 2, "login: stage 2, sid %v", sid.String())
  r2, err := lw.SendRecv(lw.MTReqChallenge(sid))
  if err != nil {
    return false, err
  }

  log("lw", 2, "login: stage 3, got response for req_challenge")
  log("lw", 2, "r2: %v", r2)

  msg, err = ParseM2Messages(r2)
  if err != nil {
    return false, err
  }

  salt := lw.MTGetSalt(msg)
  if salt == "" {
    return false, errors.New("failed to get salt from stage 3")
  }

  sl := []byte{0}
  sl = append(sl, []byte(lw.task.password)...)
  sl = append(sl, []byte(salt)...)
  d := md5.Sum(sl)
  digest := append([]byte{0}, d[:]...)

  log("lw", 2, "login: stage 3, hash %v", digest)
  r3, err := lw.SendRecv(lw.MTReqAuth(sid, lw.task.login, string(digest), string(salt)))
  if err != nil {
    return false, err
  }

  log("lw", 2, "login: stage 4, got response for req_salt")
  msg, err = ParseM2Messages(r3)
  if err != nil {
    return false, err
  }

  res, err = lw.MTGetResult(msg)
  log("lw", 2, "login: stage 5: res=%v err=%v", res, err)

  return res, err
}