package feed

import (
  "encoding/json"
  "fmt"
  // "strconv"
  "regexp"
  "time"
)

const (
  MitreFormat = iota    // MITRE data format
  DataVersion40   // Version 4.0

  OrNodeOp        // OR operator
  AndNodeOp       // And operator

  AdjacentNetwork // Adjacent Network attack vector.
  Network         // Network attack vector.
  Local           // Local attack vector.
  Physical        // Physical attack vector.

  None            // no priv req/user interaction
  Low             // low complexity/priv req
  Medium          // medium complexity/priv req
  High            // high complexity/priv req

  Required        // user interaction required

  Changed         // scope changed
  Unchanged       // scope unchanged

  Complete        // complete integrity impact
  Partial         // partial integrity impact

  Critical        // critical severity

  CvssVersion31   // CVSS version 3.1
  CvssVersion20   // CVSS version 2.0

  Single          // Single authentication
)

// TODO: parse cpe

// Data format for NVD feeds and feed items.
type DataFormat int

// Unmarshal DataFormat from JSON.
func (me *DataFormat) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "MITRE":
    *me = MitreFormat
  default:
    // return error
    return fmt.Errorf("unknown data format: %s", s)
  }

  // return success
  return nil
}

// Data version for NVD feeds and feed items.
type DataVersion int

// Unmarshal DataVersion from JSON.
func (me *DataVersion) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "4.0":
    *me = DataVersion40
  default:
    // return error
    return fmt.Errorf("unknown data version: %s", s)
  }

  // return success
  return nil
}

// partial timestamp
type PartialTime time.Time

var partialTimeRe = regexp.MustCompile("\\A\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}Z\\z")

// Unmarshal partial timestamp from JSON.
func (me *PartialTime) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // match partial string regex
  if !partialTimeRe.MatchString(s) {
    return fmt.Errorf("invalid partial time string: %s", s)
  }

  // correct string suffix
  s = s[0:16] + ":00Z"

  // unmarshal time
  var t time.Time
  if err := t.UnmarshalText([]byte(s)); err != nil {
    return err
  }

  // return success
  return nil
}

// Configuration node boolean operator.
type NodeOp int

// Unmarshal DataVersion from JSON.
func (me *NodeOp) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "AND":
    *me = AndNodeOp
  case "OR":
    *me = OrNodeOp
  default:
    // return error
    return fmt.Errorf("unknown operator: %s", s)
  }

  // return success
  return nil
}

// CVSS attack vector
type AttackVector int

// Unmarshal CVSS attack vector from JSON.
func (me *AttackVector) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "ADJACENT_NETWORK":
    *me = AdjacentNetwork
  case "LOCAL":
    *me = Local
  case "NETWORK":
    *me = Network
  case "PHYSICAL":
    *me = Physical
  default:
    // return error
    return fmt.Errorf("unknown attack vector: %s", s)
  }

  // return success
  return nil
}

// CVSS attack complexity
type AttackComplexity int

// Unmarshal CVSS attack complexity from JSON.
func (me *AttackComplexity) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "LOW":
    *me = Low
  case "MEDIUM":
    *me = Medium
  case "HIGH":
    *me = High
  default:
    // return error
    return fmt.Errorf("unknown attack complexity: %s", s)
  }

  // return success
  return nil
}

// CVSS privileges required
type PrivilegesRequired int

// Unmarshal CVSS privileges required from JSON.
func (me *PrivilegesRequired) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "NONE":
    *me = None
  case "LOW":
    *me = Low
  case "MEDIUM":
    *me = Medium
  case "HIGH":
    *me = High
  default:
    // return error
    return fmt.Errorf("unknown privileges required: %s", s)
  }

  // return success
  return nil
}

// CVSS user interaction
type UserInteraction int

// Unmarshal CVSS user interaction from JSON.
func (me *UserInteraction) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "NONE":
    *me = None
  case "REQUIRED":
    *me = Required
  default:
    // return error
    return fmt.Errorf("unknown user interaction: %s", s)
  }

  // return success
  return nil
}

// CVSS scope
type Scope int

// Unmarshal CVSS scope from JSON.
func (me *Scope) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "CHANGED":
    *me = Changed
  case "UNCHANGED":
    *me = Unchanged
  default:
    // return error
    return fmt.Errorf("unknown scope: %s", s)
  }

  // return success
  return nil
}

// CVSS integrity/availability impact level
type ImpactLevel int

// Unmarshal CVSS integrity/availability impact level from JSON.
func (me *ImpactLevel) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "NONE":
    *me = None
  case "LOW":
    *me = Low
  case "PARTIAL":
    *me = Partial
  case "HIGH":
    *me = High
  case "COMPLETE":
    *me = Complete
  default:
    // return error
    return fmt.Errorf("unknown impact level: %s", s)
  }

  // return success
  return nil
}

// CVSS score
type Score float32

// Unmarshal CVSS score from JSON.
func (me *Score) UnmarshalJSON(b []byte) error {
  // decode float, check for error
  var v float32
  if err := json.Unmarshal(b, &v); err != nil {
    return err
  }

  // check score
  if v < 0.0 || v > 10.0 {
    return fmt.Errorf("score out of bounds: %f", v)
  }

  // save result, return success
  *me = Score(v)
  return nil
}

// CVSS severity
type Severity int

// Unmarshal CVSS severity from JSON.
func (me *Severity) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "LOW":
    *me = Low
  case "MEDIUM":
    *me = Medium
  case "HIGH":
    *me = High
  case "CRITICAL":
    *me = Critical
  default:
    // return error
    return fmt.Errorf("unknown severity: %s", s)
  }

  // return success
  return nil
}

type AccessVector int

// Unmarshal CVSS V2 access vector from JSON.
func (me *AccessVector) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "ADJACENT_NETWORK":
    *me = AdjacentNetwork
  case "LOCAL":
    *me = Local
  case "NETWORK":
    *me = Network
  default:
    // return error
    return fmt.Errorf("unknown CVSS access vector: %s", s)
  }

  // return success
  return nil
}

// CVSS V2 attack complexity
type AccessComplexity int

// Unmarshal CVSS V2 access complexity from JSON.
func (me *AccessComplexity) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "LOW":
    *me = Low
  case "MEDIUM":
    *me = Medium
  case "HIGH":
    *me = High
  default:
    // return error
    return fmt.Errorf("unknown access complexity: %s", s)
  }

  // return success
  return nil
}


// CVSS V2 authentication
type Authentication int

// Unmarshal CVSS V2 authentication from JSON.
func (me *Authentication) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "NONE":
    *me = None
  case "SINGLE":
    *me = Single
  default:
    // return error
    return fmt.Errorf("unknown authentication: %s", s)
  }

  // return success
  return nil
}

// CVE metadata
type CveMetadata struct {
  // CVE ID
  Id        string `json:"ID"`

  // CVE assigner email address
  Assigner  string `json:"ASSIGNER"`
}

// CVE description string.
type Description struct {
  // Language code
  Lang  string `json:"lang"`

  // String value
  Value string `json:"value"`
}

// CVE problem type
type CveProblemType struct {
  // problem type descriptions
  Descriptions []Description `json:"description"`
}

// Slice of CVE problem types.
type CveProblemTypes struct {
  // problem types
  ProblemTypes []CveProblemType `json:"problemtype_data"`
}

// CVE reference
type CveReference struct {
  // reference URL
  Url       string    `json:"url"`

  // reference name
  Name      string    `json:"name"`

  // reference source
  RefSource string    `json:"refsource"`

  // tags
  Tags      []string  `json:"tags"`
}

// Slice of CVE references
type CveReferences struct {
  References []CveReference `json:"reference_data"`
}

// CVE item descriptions
type CveDescription struct {
  // slice of descriptions
  Descriptions  []Description `json:"description_data"`
}

// CVE data
type Cve struct {
  // feed data type
  DataType      DataType        `json:"CVE_data_type"`

  // feed data format
  DataFormat    DataFormat      `json:"CVE_data_format"`

  // feed data format version
  DataVersion   DataVersion     `json:"CVE_data_version"`

  // CVE metadata
  Metadata      CveMetadata     `json:"CVE_data_meta"`

  // CVE problem types
  ProblemTypes CveProblemTypes  `json:"problemtype"`

  // CVE references
  References   CveReferences    `json:"references"`

  // CVE description
  Description  CveDescription   `json:"description"`
}

// CPE match
type CpeMatch struct {
  // Vulnerable?
  Vulnerable bool `json:"vulnerable"`

  VersionEndExcluding string `json:"versionEndExcluding"`

  // CPE URI (FIXME: decode this)
  Cpe23Uri string `json:"cpe23Uri"`

  // CPE names (not sure if this is correct)
  Names []string `json:"cpe_name"`
}

// CVE item configuration node
type ConfigurationNode struct {
  // node operator
  Operator NodeOp `json:"operator"`

  // node children
  Children []ConfigurationNode `json:"children"`

  CpeMatches []CpeMatch `json:"cpe_match"`
}

// CVE item configurations
type ItemConfigurations struct {
  // data version
  DataVersion DataVersion `json:"CVE_data_version"`

  // slice of configuration nodes
  Nodes []ConfigurationNode `json:"nodes"`
}

// CVSS version.
type CvssV3Version int

// Unmarshal CVSS version from JSON.
func (me *CvssV3Version) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "3.1":
    *me = CvssVersion31
  default:
    // return error
    return fmt.Errorf("unknown CVSS V3 version: %s", s)
  }

  // return success
  return nil
}

// CVSS V2 version.
type CvssV2Version int

// Unmarshal CVSS version from JSON.
func (me *CvssV2Version) UnmarshalJSON(b []byte) error {
  // decode string, check for error
  var s string
  if err := json.Unmarshal(b, &s); err != nil {
    return err
  }

  // check value
  switch s {
  case "2.0":
    *me = CvssVersion20
  default:
    // return error
    return fmt.Errorf("unknown CVSS V2 version: %s", s)
  }

  // return success
  return nil
}

// CVSS V3
type CvssV3 struct {
  // CVSS V3 version
  Version CvssV3Version `json:"version"`

  // CVSS V3 vector string (FIXME: add custom type)
  VectorString string `json:"vectorString"`

  // attack vector
  AttackVector AttackVector `json:"attackVector"`

  // attack complexity
  AttackComplexity AttackComplexity `json:"attackComplexity"`

  // privileges required
  PrivilegesRequired PrivilegesRequired `json:"privilegesRequired"`

  // user interaction
  UserInteraction UserInteraction `json:"userInteraction"`

  // scope
  Scope Scope `json:"scope"`

  // integrity impact
  IntegrityImpact ImpactLevel `json:"integrityImpact"`

  // availability impact
  AvailabilityImpact ImpactLevel `json:"availabilityImpact"`

  // base score
  BaseScore Score `json:"baseScore"`

  // base severity
  BaseSeverity Severity `json:"baseSeverity"`
}

// CVSS V3 base metrics
type BaseMetricV3 struct {
  CvssV3                CvssV3  `json:"cvssV3"`
  ExploitabilityScore   Score   `json:"exploitabilityScore"`
  ImpactScore           Score   `json:"impactScore"`
}

// CVSS V2
type CvssV2 struct {
  // CVSS V2 version
  Version CvssV2Version `json:"version"`

  // CVSS V3 vector string (FIXME: add custom type)
  VectorString string `json:"vectorString"`

  // attack vector
  AccessVector AccessVector `json:"accessVector"`

  // attack complexity
  AccessComplexity AccessComplexity `json:"accessComplexity"`

  // authentication
  Authentication Authentication `json:"authentication"`

  ConfidentialityImpact ImpactLevel `json:"confidentialityImpact"`
  IntegrityImpact ImpactLevel `json:"integrityImpact"`
  AvailabilityImpact ImpactLevel `json:"availabilityImpact"`

  // base score
  BaseScore Score `json:"baseScore"`
}

// CVSS V2 base metrics
type BaseMetricV2 struct {
  CvssV2                  CvssV2    `json:"cvssV2"`
  severity                Severity  `json:"severity"`
  ExploitabilityScore     Score     `json:"impactScore"`
  ImpactScore             Score     `json:"impactScore"`
  InsufficientInfo        bool      `json:"acInsufInfo"`
  ObtainAllPrivilege      bool      `json:"obtainAllPrivilege"`
  ObtainUserPrivilege     bool      `json:"obtainUserPrivilege"`
  ObtainOtherPrivilege    bool      `json:"obtainOtherPrivilege"`
  UserInteractionRequired bool      `json:"userInteractionRequired"`
}

// Item impact
type Impact struct {
  // CVSS V3 base metrics
  BaseMetricV3 BaseMetricV3 `json:"baseMetricV3"`

  // CVSS V2 base metrics
  BaseMetricV2 BaseMetricV2 `json:"baseMetricV2"`
}

// CVE feed item
type Item struct {
  // item CVE data
  Cve Cve `json:"cve"`

  // item configuration
  Configurations ItemConfigurations `json:"configurations"`

  // item impact
  Impact Impact `json:"impact"`

  // item published date
  PublishedDate PartialTime `json:"publishedDate"`

  // last modification date
  LastModifiedDate PartialTime `json:"lastModifiedDate"`
}

// NVD feed
type Feed struct {
  // feed data type
  DataType      DataType    `json:"CVE_data_type"`

  // feed data format
  DataFormat    DataFormat  `json:"CVE_data_format"`

  // feed data format version
  DataVersion   DataVersion `json:"CVE_data_version"`

  // number of CVEs in feed
  NumCVEs       uint64      `json:"CVE_data_numberOfCVEs,string"`

  // data timestamp
  Timestamp     PartialTime `json:"CVE_data_timestamp"`

  // CVE items
  Items         []Item      `json:"CVE_Items"`
}